RUN echo $CACHEBUSTER

# Setup build and environment
COPY devel.usertools /usertools
COPY devel.bashrc /root/.bashrc
COPY ld.so.conf /dt10/etc/

# Setup JAX Python environment.
FROM devel as jax
RUN /setup.packages.sh /cuda.packages.txt
COPY jax.requirements.txt /devel.requirements.txt
COPY setup.python.sh /setup.python.sh
RUN /setup.python.sh python3.9 devel.requirements.txt

# Test dependencies for pip tests
grpcio ~= 1.42.0
portpicker ~= 1.5.2
scipy ~= 1.7.3
jax ~= 0.2.26

    #   h5py
    #   jax
    #   keras-nightly
    #   ml-dtypes
    #   opt-einsum
    #   scipy
    #   tb-nightly
opt-einsum==3.3.0 \
    --hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
    --hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549
    # via
    #   -r requirements.in
    #   jax
packaging==23.2 \

        --lockfile=requirements.txt:./requirements_lock_3_12.txt
        --lockfile=requirements.txt:./ci/official/containers/linux_arm64/devel.requirements.txt
        --lockfile=requirements.txt:./ci/official/containers/linux_arm64/jax.requirements.txt

reachable and exploitable through production-grade, benign models.

### Compilation

Compiling models via the recommended entry points described in
[XLA](https://www.tensorflow.org/xla) and
[JAX](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html)
documentation should be safe, while some of the testing and debugging tools that
come with the compiler are not designed to be used with untrusted data and

keras-nightly ~= 3.0.0.dev
tb-nightly ~= 2.17.0.a

# Test dependencies
grpcio >= 1.24.3, < 2.0
portpicker == 1.6.0
scipy == 1.11.3
requests == 2.31.0
packaging==23.2
setuptools==68.2.2

    #   h5py
    #   jax
    #   keras-nightly
    #   ml-dtypes
    #   opt-einsum
    #   scipy
    #   tb-nightly
opt-einsum==3.3.0 \
    --hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
    --hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549
    # via
    #   -r requirements.in
    #   jax
packaging==23.2 \

        "//tensorflow/compiler/mlir/lite/python:flatbuffer_to_mlir",
        "//tensorflow/compiler/mlir/lite/python:graphdef_to_tfl_flatbuffer",
        "//tensorflow/compiler/mlir/lite/python:jax_to_tfl_flatbuffer",
        "//tensorflow/compiler/mlir/lite/python:saved_model_to_tfl_flatbuffer",
        "//tensorflow/compiler/mlir/lite/quantization/lite:quantize_model",

        tf.constant(2, dtype=tf.int32) b = tf.constant(3, dtype=tf.int32)
        result = client.multiply(a, b)`

*   `tf.lite`:

    *   Add experimental API `experimental_from_jax` to support conversion from
        Jax models to TensorFlow Lite.
    *   Support uint32 data type for cast op.
    *   Support int8 data type for cast op.
    *   Add experimental quantization debugger `tf.lite.QuantizationDebugger`

    assignees.push(
        'DEKHTIARJonathan', 'meena-at-work', 'nluehr', 'pjannaty', 'poulsbo');
  } else if (
      lowercased_title.includes('nvidia') && domain.includes('nvidia.com')) {
    if (lowercased_title.includes('jax')) {
      assignees.push('hawkinsp', 'yashk2810', 'skye');
    }
    if (lowercased_title.includes('xla') || lowercased_title.includes('gpu')) {
      assignees.push('cheshire', 'gcforster', 'reedwm', 'chsigg', 'xla-rotation');