test tree_util changes

.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: |