Jax grad doesn't work with tfp jax module's Gram-Schmidt function because of a lax for loop
@MintyMH One workaround could be to rewrite the Gram-Schmidt function using vectorized operations and functional programming constructs like jax.map and jax.lax.cond. This avoids control flow and ensures compatibility with grad.
Thank you! This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62828">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62828">No</a>

### Issue type

Support

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

0.24.0.dev20240121

### Custom code

Yes

### OS platform and distribution

Windows 10 Home 10.0.19045 Build 19045

### Mobile device

_No response_

### Python version

3.11.4

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

The gram-schmidt function seems to not work with jax's grad function because of a problematic lax for loop.

### Standalone code to reproduce the issue

```shell

import numpy as np
from tensorflow_probability.substrates import jax as tfp

import jax
import jax.numpy as jnp
from jax import grad, jit, vmap

from jax import config
config.update("jax_enable_x64", True)

matrices = np.random.random((20, 20, 4, 4))

def my_function(matrices):
    Q = tfp.math.gram_schmidt(matrices)
    return jnp.linalg.norm(Q)

temp = my_function(matrices) # runs fine

my_grad = jit(grad(my_function))
 
grad_temp = my_grad(matrices)
```


### Relevant log output

```shell
"name": "ValueError",
	"message": "Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop.",
	"stack": "---------------------------------------------------------------------------
JaxStackTraceBeforeTransformation         Traceback (most recent call last)
File <frozen runpy>:198, in _run_module_as_main()

File <frozen runpy>:88, in _run_code()

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\ipykernel_launcher.py:17
     15 from ipykernel import kernelapp as app
---> 17 app.launch_new_instance()

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\traitlets\\config\\application.py:992, in launch_instance()
    991 app.initialize(argv)
--> 992 app.start()

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\ipykernel\\kernelapp.py:711, in start()
    710 try:
--> 711     self.io_loop.start()
    712 except KeyboardInterrupt:

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\tornado\\platform\\asyncio.py:195, in start()
    194 def start(self) -> None:
--> 195     self.asyncio_loop.run_forever()

File c:\\Users\\bjmon\\anaconda3\\Lib\\asyncio\\base_events.py:607, in run_forever()
    606 while True:
--> 607     self._run_once()
    608     if self._stopping:

File c:\\Users\\bjmon\\anaconda3\\Lib\\asyncio\\base_events.py:1922, in _run_once()
   1921     else:
-> 1922         handle._run()
   1923 handle = None

File c:\\Users\\bjmon\\anaconda3\\Lib\\asyncio\\events.py:80, in _run()
     79 try:
---> 80     self._context.run(self._callback, *self._args)
     81 except (SystemExit, KeyboardInterrupt):

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\ipykernel\\kernelbase.py:510, in dispatch_queue()
    509 try:
--> 510     await self.process_one()
    511 except Exception:

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\ipykernel\\kernelbase.py:499, in process_one()
    498         return None
--> 499 await dispatch(*args)

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\ipykernel\\kernelbase.py:406, in dispatch_shell()
    405     if inspect.isawaitable(result):
--> 406         await result
    407 except Exception:

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\ipykernel\\kernelbase.py:729, in execute_request()
    728 if inspect.isawaitable(reply_content):
--> 729     reply_content = await reply_content
    731 # Flush output before sending the reply.

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\ipykernel\\ipkernel.py:411, in do_execute()
    410 if with_cell_id:
--> 411     res = shell.run_cell(
    412         code,
    413         store_history=store_history,
    414         silent=silent,
    415         cell_id=cell_id,
    416     )
    417 else:

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\ipykernel\\zmqshell.py:531, in run_cell()
    530 self._last_traceback = None
--> 531 return super().run_cell(*args, **kwargs)

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\IPython\\core\\interactiveshell.py:3006, in run_cell()
   3005 try:
-> 3006     result = self._run_cell(
   3007         raw_cell, store_history, silent, shell_futures, cell_id
   3008     )
   3009 finally:

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\IPython\\core\\interactiveshell.py:3061, in _run_cell()
   3060 try:
-> 3061     result = runner(coro)
   3062 except BaseException as e:

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\IPython\\core\\async_helpers.py:129, in _pseudo_sync_runner()
    128 try:
--> 129     coro.send(None)
    130 except StopIteration as exc:

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\IPython\\core\\interactiveshell.py:3266, in run_cell_async()
   3263 interactivity = \"none\" if silent else self.ast_node_interactivity
-> 3266 has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
   3267        interactivity=interactivity, compiler=compiler, result=result)
   3269 self.last_execution_succeeded = not has_raised

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\IPython\\core\\interactiveshell.py:3445, in run_ast_nodes()
   3444     asy = compare(code)
-> 3445 if await self.run_code(code, result, async_=asy):
   3446     return True

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\IPython\\core\\interactiveshell.py:3505, in run_code()
   3504     else:
-> 3505         exec(code_obj, self.user_global_ns, self.user_ns)
   3506 finally:
   3507     # Reset our crash handler in place

Cell In[2], line 15
     13 my_grad = jit(grad(my_function))
---> 15 grad_temp = my_grad(matrices)

Cell In[2], line 5, in my_function()
      3 def my_function(matrices):
----> 5     Q = tfp.math.gram_schmidt(matrices)
      7     return jnp.linalg.norm(Q)

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\tensorflow_probability\\substrates\\jax\\math\\gram_schmidt.py:83, in gram_schmidt()
     81   return vecs, i + 1
---> 83 vectors, _ = tf.while_loop(cond, body_fn, (vectors, tf.zeros([], tf.int32)))
     84 vec_norm = tf.linalg.norm(vectors, ord=2, axis=-2, keepdims=True)

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\tensorflow_probability\\python\\internal\\backend\\jax\\control_flow.py:90, in _while_loop_jax()
     89     return cond(*args)
---> 90   return lax.while_loop(override_cond_fn, override_body_fn, loop_vars)
     91 elif back_prop:

JaxStackTraceBeforeTransformation: ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop.

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

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

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
Cell In[2], line 15
     11 print(temp)
     13 my_grad = jit(grad(my_function))
---> 15 grad_temp = my_grad(matrices)

    [... skipping hidden 21 frame]

File c:\\Users\\bjmon\\anaconda3\\Lib\\site-packages\\jax\\_src\\lax\\control_flow\\loops.py:1549, in _while_transpose_error(*_, **kwargs)
   1548 def _while_transpose_error(*_, **kwargs):
-> 1549   raise ValueError(\"Reverse-mode differentiation does not work for \"
   1550                    \"lax.while_loop or lax.fori_loop with dynamic start/stop values. \"
   1551                    \"Try using lax.scan, or using fori_loop with static start/stop.\")

ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop."
```
