01-Optimization MPAX calling error

[evelynmitchell]

01-Optimization notebook In Predefined MPAX:

ERROR:2025-03-19 17:47:03,718:jax._src.xla_bridge:629: Jax plugin configuration error: Exception when calling jax_plugins.xla_cuda12.initialize()
Traceback (most recent call last):
  File "/usr/local/lib/python3.11/dist-packages/jax/_src/xla_bridge.py", line 627, in discover_pjrt_plugins
    plugin_module.initialize()
  File "/usr/local/lib/python3.11/dist-packages/jax_plugins/xla_cuda12/__init__.py", line 103, in initialize
    triton.register_compilation_handler(
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'jax._src.lib.triton' has no attribute 'register_compilation_handler'
ERROR:jax._src.xla_bridge:Jax plugin configuration error: Exception when calling jax_plugins.xla_cuda12.initialize()
Traceback (most recent call last):
  File "/usr/local/lib/python3.11/dist-packages/jax/_src/xla_bridge.py", line 627, in discover_pjrt_plugins
    plugin_module.initialize()
  File "/usr/local/lib/python3.11/dist-packages/jax_plugins/xla_cuda12/__init__.py", line 103, in initialize
    triton.register_compilation_handler(
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'jax._src.lib.triton' has no attribute 'register_compilation_handler'

---------------------------------------------------------------------------

AttributeError                            Traceback (most recent call last)

[<ipython-input-12-363ea713f49a>](https://localhost:8080/#) in <cell line: 0>()
      1 # solve model
      2 optmodel = shortestPathModel(grid=(5,5)) # init model
----> 3 optmodel.setObj(cost) # set objective function
      4 sol, obj = optmodel.solve() # solve
      5 # print res

[/usr/local/lib/python3.11/dist-packages/pyepo/model/mpax/mpaxmodel.py](https://localhost:8080/#) in setObj(self, c)
    100             c (np.ndarray / list): cost of objective function
    101         """
--> 102         if c.shape[-1] != self.num_cost:
    103             raise ValueError("Size of cost vector cannot match vars.")
    104         # check if c is a PyTorch tensor

AttributeError: 'list' object has no attribute 'shape'

Solution is

# solve model
optmodel = shortestPathModel(grid=(5,5)) # init model
optmodel.setObj(np.array(cost)) # set objective function, converting cost to a NumPy array
sol, obj = optmodel.solve() # solve
# print res
print('Obj: {}'.format(obj))
for i, e in enumerate(optmodel.arcs):
    if sol[i] > 1e-3:
        print(e)

[LucasBoTang]

Hi, there was a small bug earlier — the setObj function expects a NumPy array instead of a plain list. This has now been fixed.