test tree_util changes (#1396)

This PR is for testing my changes to AXLearn, to avoid having this
error:
```python 
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]   File "/opt/axlearn/axlearn/common/trainer.py", line 1162, in compile_train_step
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]     lowered_train_step = jit_train_step.lower(trainer_state, input_batch)
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]   File "/opt/axlearn/axlearn/common/trainer.py", line 1204, in _train_step
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]     fwd_bwd_outputs, learner_output_collection = F(
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]   File "/opt/axlearn/axlearn/common/module.py", line 1096, in functional
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]     method_outputs, output_collection = fn(*args, **kwargs)
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]   File "/opt/axlearn/axlearn/common/module.py", line 1032, in __call__
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]     raise_for_cycles(dict(context=self.context, args=args, kwargs=kwargs))
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]   File "/opt/axlearn/axlearn/common/utils.py", line 1930, in raise_for_cycles
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]     cycles = find_cycles(tree)
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]   File "/opt/axlearn/axlearn/common/utils.py", line 1924, in find_cycles
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]     return _find_cycles(tree, key_path=[], seen=[])
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]   File "/opt/axlearn/axlearn/common/utils.py", line 1914, in _find_cycles
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]     items = pytree_children(tree)
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]   File "/opt/axlearn/axlearn/common/utils.py", line 1872, in pytree_children
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model]     registry_with_keypaths = jax._src.tree_util._registry_with_keypaths
[pod/axlearn-fuji-3b-14466154556-kgkm5/axlearn-fuji-model] AttributeError: module 'jax._src.tree_util' has no attribute '_registry_with_keypaths'
```
with recent JAX version. 
The changes are on [my AXLearn
repo](https://github.com/Steboss/axlearn/blob/sbosisio/tree_util/axlearn/common/utils.py#L1862)

I'll flag this error to AXLearn team

---------

Co-authored-by: Olli Lupton <[email protected]>

Author: Steboss

.github/container/Dockerfile.axlearn

@@ -1,6 +1,6 @@
 # syntax=docker/dockerfile:1-labs
 ARG BASE_IMAGE=ghcr.io/nvidia/jax-mealkit:jax
-ARG URLREF_AXLEARN=https://github.com/Steboss/axlearn.git#main
+ARG URLREF_AXLEARN=https://github.com/apple/axlearn.git
 ARG SRC_PATH_AXLEARN=/opt/axlearn
 
 ###############################################################################

.github/container/Dockerfile.jax

@@ -53,7 +53,7 @@ RUN build-jax.sh \
     --src-path-jax ${SRC_PATH_JAX} \
     --src-path-xla ${SRC_PATH_XLA} \
     --sm all \
-    --xla-arm64-patch /opt/xla-arm64-neon.patch \ 
+    --xla-arm64-patch /opt/xla-arm64-neon.patch \
     --clean
 
 ## Transformer engine: check out source and build wheel

.github/container/manifest.yaml

@@ -101,3 +101,7 @@ pathwaysutils:
   tracking_ref: main
   latest_verified_commit: 359776d454940ffaa337c36d1df16308d44a95a9
   mode: pip-vcs
+axlearn:
+  url: https://github.com/Steboss/axlearn.git
+  tracking_ref: sbosisio/working_branch
+  mode: git-clone

.github/workflows/_ci.yaml

@@ -134,6 +134,8 @@ jobs:
       CONTAINER_NAME: axlearn
       DOCKERFILE: .github/container/Dockerfile.axlearn
       RUNNER_SIZE: large
+      EXTRA_BUILD_ARGS: |
+        URLREF_AXLEARN=${{ fromJson(inputs.SOURCE_URLREFS).AXLEARN }}
     secrets: inherit
 
   collect-docker-tags:
@@ -438,6 +440,7 @@ jobs:
   #     ARTIFACTS: |
   #       test-equinox.log
   #   secrets: inherit
+
   test-te-h100:
     needs: build-jax
     if: inputs.ARCHITECTURE == 'amd64'
@@ -478,7 +481,7 @@ jobs:
           # merge the log files
           cat \
             /log/pytest-report-L0-unittest.jsonl
-            /log/pytest-report-L0-distributed-unittest.jsonl 
+            /log/pytest-report-L0-distributed-unittest.jsonl
             > /log/pytest-report.jsonl
         EOF
       STATISTICS_SCRIPT: |