# Path: jax/numerics.py
# def lerp_weight(x, xs):
#   """Linear interpolation weight from a sample at x to xs.
# 
#   Returns the linear interpolation weight of a "query point" at coordinate `x`
#   with respect to a "sample" at coordinate `xs`.
# 
#   The integer coordinates `x` are at pixel centers.
#   The floating point coordinates `xs` are at pixel edges.
#   (OpenGL convention).
# 
#   Args:
#     x: "Query" point position.
#     xs: "Sample" position.
# 
#   Returns:
#     - 1 when x = xs.
#     - 0 when |x - xs| > 1.
#   """
#   dx = x - xs
#   abs_dx = abs(dx)
#   return jnp.maximum(1.0 - abs_dx, 0.0)
# 
# Path: jax/numerics.py
# def smoothed_lerp_weight(x, xs, eps=1e-8):
#   """Smoothed version of `LerpWeight` with gradients more suitable for backprop.
# 
#   Let f(x, xs) = LerpWeight(x, xs)
#                = max(1 - |x - xs|, 0)
#                = max(1 - |dx|, 0)
# 
#   f is not smooth when:
#   - |dx| is close to 0. We smooth this by replacing |dx| with
#     SmoothedAbs(dx, eps) = sqrt(dx * dx + eps), which has derivative
#     dx / sqrt(dx * dx + eps).
#   - |dx| = 1. When smoothed, this happens when dx = sqrt(1 - eps). Like ReLU,
#     We just ignore this (in the implementation below, when the floats are
#     exactly equal, we choose the SmoothedAbsGrad path since it is more useful
#     than returning a 0 gradient).
# 
#   Args:
#     x: "Query" point position.
#     xs: "Sample" position.
#     eps: a small number.
# 
#   Returns:
#     max(1 - |dx|, 0) where |dx| is smoothed_abs(dx).
#   """
#   dx = x - xs
#   abs_dx = smoothed_abs(dx, eps)
#   return jnp.maximum(1.0 - abs_dx, 0.0)
# 
# Path: jax/numerics.py
# def smoothed_lerp_weight_grad(x, xs, eps=1e-8):
#   """Gradient of smoothed_lerp_weight."""
#   dx = x - xs
#   abs_dx = smoothed_abs(dx, eps)
#   grad = smoothed_abs_grad(dx, eps)
#   return jnp.where(abs_dx > 1, 0, grad)


import jax
import jax.numpy as jnp
from .numerics import lerp_weight
from .numerics import smoothed_lerp_weight
from .numerics import smoothed_lerp_weight_grad

  Returns:
    The vector-Jacobian product codomain_tangent * J^T(grid, guide) with shape
    (h, w) (the same as that of the primal `guide`).

  """
  ii, jj = jnp.meshgrid(
      jnp.arange(guide.shape[0]), jnp.arange(guide.shape[1]), indexing='ij')

  scale_i = grid.shape[0] / guide.shape[0]
  scale_j = grid.shape[1] / guide.shape[1]
  grid_depth = grid.shape[2]

  gif = (ii + 0.5) * scale_i
  gjf = (jj + 0.5) * scale_j
  gkf = guide * grid.shape[2]

  # Compute trilinear interpolation weights without clamping.
  gi0 = jnp.floor(gif - 0.5).astype(jnp.int32)
  gj0 = jnp.floor(gjf - 0.5).astype(jnp.int32)
  gk0 = jnp.floor(gkf - 0.5).astype(jnp.int32)
  gi1 = gi0 + 1
  gj1 = gj0 + 1
  gk1 = gk0 + 1

  wi0 = lerp_weight(gi0 + 0.5, gif)
  wi1 = lerp_weight(gi1 + 0.5, gif)
  wj0 = lerp_weight(gj0 + 0.5, gjf)
  wj1 = lerp_weight(gj1 + 0.5, gjf)
  dwk0 = grid_depth * smoothed_lerp_weight_grad(gk0 + 0.5, gkf)

# Path: jax/numerics.py
# def lerp_weight(x, xs):
#   """Linear interpolation weight from a sample at x to xs.
# 
#   Returns the linear interpolation weight of a "query point" at coordinate `x`
#   with respect to a "sample" at coordinate `xs`.
# 
#   The integer coordinates `x` are at pixel centers.
#   The floating point coordinates `xs` are at pixel edges.
#   (OpenGL convention).
# 
#   Args:
#     x: "Query" point position.
#     xs: "Sample" position.
# 
#   Returns:
#     - 1 when x = xs.
#     - 0 when |x - xs| > 1.
#   """
#   dx = x - xs
#   abs_dx = abs(dx)
#   return jnp.maximum(1.0 - abs_dx, 0.0)
# 
# Path: jax/numerics.py
# def smoothed_lerp_weight(x, xs, eps=1e-8):
#   """Smoothed version of `LerpWeight` with gradients more suitable for backprop.
# 
#   Let f(x, xs) = LerpWeight(x, xs)
#                = max(1 - |x - xs|, 0)
#                = max(1 - |dx|, 0)
# 
#   f is not smooth when:
#   - |dx| is close to 0. We smooth this by replacing |dx| with
#     SmoothedAbs(dx, eps) = sqrt(dx * dx + eps), which has derivative
#     dx / sqrt(dx * dx + eps).
#   - |dx| = 1. When smoothed, this happens when dx = sqrt(1 - eps). Like ReLU,
#     We just ignore this (in the implementation below, when the floats are
#     exactly equal, we choose the SmoothedAbsGrad path since it is more useful
#     than returning a 0 gradient).
# 
#   Args:
#     x: "Query" point position.
#     xs: "Sample" position.
#     eps: a small number.
# 
#   Returns:
#     max(1 - |dx|, 0) where |dx| is smoothed_abs(dx).
#   """
#   dx = x - xs
#   abs_dx = smoothed_abs(dx, eps)
#   return jnp.maximum(1.0 - abs_dx, 0.0)
# 
# Path: jax/numerics.py
# def smoothed_lerp_weight_grad(x, xs, eps=1e-8):
#   """Gradient of smoothed_lerp_weight."""
#   dx = x - xs
#   abs_dx = smoothed_abs(dx, eps)
#   grad = smoothed_abs_grad(dx, eps)
#   return jnp.where(abs_dx > 1, 0, grad)


# Path: jax/bilateral_slice.py
import jax
import jax.numpy as jnp
from .numerics import lerp_weight
from .numerics import smoothed_lerp_weight
from .numerics import smoothed_lerp_weight_grad


  Returns:
    The vector-Jacobian product codomain_tangent * J^T(grid, guide) with shape
    (h, w) (the same as that of the primal `guide`).

  """
  ii, jj = jnp.meshgrid(
      jnp.arange(guide.shape[0]), jnp.arange(guide.shape[1]), indexing='ij')

  scale_i = grid.shape[0] / guide.shape[0]
  scale_j = grid.shape[1] / guide.shape[1]
  grid_depth = grid.shape[2]

  gif = (ii + 0.5) * scale_i
  gjf = (jj + 0.5) * scale_j
  gkf = guide * grid.shape[2]

  # Compute trilinear interpolation weights without clamping.
  gi0 = jnp.floor(gif - 0.5).astype(jnp.int32)
  gj0 = jnp.floor(gjf - 0.5).astype(jnp.int32)
  gk0 = jnp.floor(gkf - 0.5).astype(jnp.int32)
  gi1 = gi0 + 1
  gj1 = gj0 + 1
  gk1 = gk0 + 1

  wi0 = lerp_weight(gi0 + 0.5, gif)
  wi1 = lerp_weight(gi1 + 0.5, gif)
  wj0 = lerp_weight(gj0 + 0.5, gjf)
  wj1 = lerp_weight(gj1 + 0.5, gjf)
  dwk0 = grid_depth * smoothed_lerp_weight_grad(gk0 + 0.5, gkf)

  dwk1 = grid_depth * smoothed_lerp_weight_grad(gk1 + 0.5, gkf)