Run all tests using different modules

[NeilGirdhar]

Run using:

• numpy,
• jax.numpy, and
• array_api_strict

Steps

• Hardest part: Find substitutes for grad, vmap, jacobian, hessian, stop_gradient, jacfwd, etc.
• Make a fixture xp: ModuleType..
• Use the fixture to create all arrays in the tests.
• Deal with the problem of jax.random. Need to switch on self.array_namespace() and provide appropriate samples.
• Deal with the problem of jit. Need to switch and only jit if self demands it?

def jit_if_ok(f) -> ...:
  def new_f(self, ...):
      if needs_jit(self):
        return jit(f)(...)
      return f(...)
  return new_f

• In testing, we need assert_tree_allclose, and one use of tree.map.