aws
[feature-request] Support for JAX container
Concise Description: I'd like to use JAX for distributed training of LLMs. In addition, the new release of Keras supports JAX as a backend in addition to TF.
Describe the solution you'd like I'd like either a separate JAX container or jaxlib included in the TF container since the TF ecosystem (data loading, serving, etc) supports JAX.
Describe alternatives you've considered I could install JAX on top of the TF container.
@sbhavani Thank you for submitting the feature request. We will discuss and prioritize the adoption of JAX into the DLC. Unfortunately I cannot comment on a timeline for this feature. We will update the public documentation and github release notes once the feature is released.
