Error using jit for speed up

[PoulomiPradhan]

I have a function where I am applying grad and getting values. But when I am using jit for speedup, I am getting Jax error.

gradObjFunc = grad(objective_function)
p = (np.random.rand(5))/100000000
gradObjFunc(p) 

This works fine but the code below throws error

jit_gradObjFunc = jit(grad(objective_function))
jit_gradObjFunc(p)

Error :

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

ConcretizationTypeError                   Traceback (most recent call last)

Cell In [124], line 1
----> 1 jit_gradObjFunc(p)

    [... skipping hidden 24 frame]

Cell In [106], line 3, in objective_function(p)
      1 def objective_function(p):
      2     # the objective function calculates the loss give lens with shape parameter p
----> 3     sm = lens(f, R1, R2, d_L, wv, n, pupil_r, p)
      5     thi = sm.gaps[0].thi
      6     number_of_colors = 5

Cell In [100], line 2, in lens(f, R1, R2, d_L, wv, n, pupil_r, p)
      1 def lens(f, R1, R2, d_L, wv, n, pupil_r, p = [0]):
----> 2     opm = OpticalModel()
      3     sm  = opm['seq_model']
      4     osp = opm['optical_spec']

File ~/Downloads/rayoptics/optical/opticalmodel.py:134, in OpticalModel.__init__(self, radius_mode, specsheet, **kwargs)
    130     self.set_from_specsheet()
    132 if kwargs.get('do_init', True):
    133     # need to do this after OpticalSpec is initialized
--> 134     self.seq_model.update_model()
    135     elements_from_sequence(self.ele_model,
    136                            self.seq_model,
    137                            self.part_tree)

File ~/Downloads/rayoptics/seq/sequential.py:454, in SequentialModel.update_model(self, **kwargs)
    452 if g is not None:
    453     n_after = self.rndx[i][ref_wl]
--> 454     if z_dir_after < 0:
    455         n_after = -n_after
    456     ifc.delta_n = n_after - n_before

    [... skipping hidden 1 frame]

File ~/.local/lib/python3.10/site-packages/jax/core.py:1189, in concretization_function_error.<locals>.error(self, arg)
   1188 def error(self, arg):
-> 1189   raise ConcretizationTypeError(arg, fname_context)

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
The problem arose with the `bool` function. 
The error occurred while tracing the function objective_function at /tmp/ipykernel_1767043/2661071087.py:1 for jit. This value became a tracer due to JAX operations on these lines:

  operation a:f32[] = xla_call[
  call_jaxpr={ lambda ; b:i32[] c:i32[]. let
      d:f32[] = convert_element_type[new_dtype=float32 weak_type=True] b
      e:f32[] = convert_element_type[new_dtype=float32 weak_type=True] c
      f:bool[] = xla_call[
        call_jaxpr={ lambda ; g:f32[]. let
            h:i32[] = bitcast_convert_type[new_dtype=int32] g
            i:i32[] = shift_right_arithmetic h 31
            j:bool[] = convert_element_type[new_dtype=bool weak_type=False] i
          in (j,) }
        name=signbit
      ] e
      k:f32[] = abs d
      l:f32[] = neg k
      m:f32[] = abs d
      n:f32[] = xla_call[
        call_jaxpr={ lambda ; o:bool[] p:f32[] q:f32[]. let
            r:f32[] = select_n o q p
          in (r,) }
        name=_where
      ] f l m
    in (n,) }
  name=copysign
] s s
    from line /home/poulomi/Downloads/rayoptics/seq/sequential.py:441 (update_model)

  operation a:f32[] = convert_element_type[new_dtype=float32 weak_type=True] b
    from line /home/poulomi/Downloads/rayoptics/seq/sequential.py:443 (update_model)

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError

Can someone kindly help here ?

[jakevdp]

jit is not compatible with functions that use Python control flow operators on traced values; so this line in your function is problematic:

    453     n_after = self.rndx[i][ref_wl]
--> 454     if z_dir_after < 0:
    455         n_after = -n_after
    456     ifc.delta_n = n_after - n_before

There is some general discussion of this here: https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#control-flow.

If you share a short reproduction of the issue here, we could give more pointed advice for your particular situation, but assuming lines 454-455 are the only problem, you may be able to replace them with this:

n_after *= jnp.sign(z_dir_after) 

[PoulomiPradhan]

Thank you for your reply. I changed the code as in line no 454

n_after *= jnp.sign(z_dir_after)

now it's throwing error from somewhere else.

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

ConcretizationTypeError                   Traceback (most recent call last)

Cell In [18], line 1
----> 1 jit_gradObjFunc(p)

    [... skipping hidden 24 frame]

Cell In [13], line 3, in objective_function(p)
      1 def objective_function(p):
      2     # the objective function calculates the loss give lens with shape parameter p
----> 3     sm = lens(f, R1, R2, d_L, wv, n, pupil_r, p)
      5     thi = sm.gaps[0].thi
      6     number_of_colors = 5

Cell In [4], line 28, in lens(f, R1, R2, d_L, wv, n, pupil_r, p)
     25 sm.add_surface([R2, 0])
     27 sm.gaps[-1].thi = remaining_height - d_L
---> 28 opm.update_model()
     29 return sm

File ~/Downloads/rayoptics/optical/opticalmodel.py:331, in OpticalModel.update_model(self, **kwargs)
    317 def update_model(self, **kwargs):
    318     """Model and its constituents are updated.
    319 
    320     Args:
   (...)
    329 
    330     """
--> 331     self['seq_model'].update_model(**kwargs)
    332     self['optical_spec'].update_model(**kwargs)
    333     self.update_optical_properties(**kwargs)

File ~/Downloads/rayoptics/seq/sequential.py:464, in SequentialModel.update_model(self, **kwargs)
    461         self.z_dir[i] = z_dir_after
    463     # call update() on the surface interface
--> 464     ifc.update()
    466 self.gbl_tfrms = self.compute_global_coords()
    467 self.lcl_tfrms = self.compute_local_transforms()

File ~/Downloads/rayoptics/elem/surface.py:122, in Surface.update(self)
    120 def update(self):
    121     super().update()
--> 122     self.profile.update()

File ~/Downloads/rayoptics/elem/profiles.py:947, in RadialPolynomial.update(self)
    946 def update(self):
--> 947     self.gen_coef_list()
    948     return self

File ~/Downloads/rayoptics/elem/profiles.py:924, in RadialPolynomial.gen_coef_list(self)
    922     self.max_nonzero_coef = -1
    923 for i, c in enumerate(self.coefs):
--> 924     if c != 0.0:
    925         self.max_nonzero_coef = i
    926 self.max_nonzero_coef += 1
    [... skipping hidden 1 frame]

File ~/.local/lib/python3.10/site-packages/jax/core.py:1189, in concretization_function_error.<locals>.error(self, arg)
   1188 def error(self, arg):
-> 1189   raise ConcretizationTypeError(arg, fname_context)

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=0/1)>
The problem arose with the `bool` function. 
The error occurred while tracing the function objective_function at /tmp/ipykernel_2272287/2661071087.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument 'p'.

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError

Can you please let me know how should I fix this line of code in line 924 ? I think fixing one by one line in jax wherever I am getting this control flow issue, is the only way to solve this problem.

Thank you very much for your response.

[jakevdp]

The issue is the same as above: you cannot do value-dependent Python control flow within JIT. You'll have to update your code to avoid these kinds of if statements. If you're able to share a minimal reproducible example, I can probably give more guidance, but it's difficult to guess from the traceback alone what your code does and how it could be reworked.