def inside_conical_frustum(x, d, t0, t1, r, ttol=1e-6, rtol=1e-6):
  """Test if `x` is inside the conical frustum specified by the other inputs."""
  d_normsq = jnp.sum(d**2)
  d_norm = jnp.sqrt(d_normsq)
  x_normsq = jnp.sum(x**2, -1)
  x_norm = jnp.sqrt(x_normsq)
  xd = math.matmul(x, d)
  is_inside = (
      (t0 - ttol) <= xd / d_normsq) & (xd / d_normsq <= (t1 + ttol)) & (
          (xd / (d_norm * x_norm)) >=
          (1 / jnp.sqrt(1 + r**2 / d_normsq) - rtol))
  return is_inside
Test if `x` is inside the conical frustum specified by the other inputs.
internal/mip_test.py