Update to jax/jaxlib 0.4.2 (#48)

There were some fixes required to update to jax 0.4.2. This includes a change to GlobalOp initialization.
iree-org · Jan 26, 2023 · 26006ef · 26006ef
1 parent fa3e3d9
commit 26006ef
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/iree/jax/ir_utils.py b/iree/jax/ir_utils.py
@@ -43,11 +43,11 @@ def create_global(symbol_table: ir.SymbolTable,
                   visibility: str = "private",
                   initial_value: Optional[ir.Attribute] = None) -> str:
   op = ml_program_d.GlobalOp(
-      sym_visibility=ir.StringAttr.get(visibility),
-      sym_name=ir.StringAttr.get(symbol),
-      type=ir.TypeAttr.get(ir_type),
+      ir.StringAttr.get(symbol),
+      ir.TypeAttr.get(ir_type),
       is_mutable=ir.UnitAttr.get() if mutable else None,
       value=initial_value,
+      sym_visibility=ir.StringAttr.get(visibility),
   )
   symbol_table.insert(op)
   # Must get the symbol name after insert, since it may be renamed.

diff --git a/version_info.json b/version_info.json
@@ -3,7 +3,7 @@
     "iree-compiler": "20221121.334",
     "iree-runtime": "20221121.334",
     "iree-tools-xla": "20221121.334",
-    "jax": "0.4.1",
-    "jaxlib": "0.4.1"
+    "jax": "0.4.2",
+    "jaxlib": "0.4.2"
   }
 }