jax_sn.operators.transport_operator.implementations

Concrete strategy implementations for the transport operator.

These modules hold the bulk of the sweep/forward kernels; the public symbols are re-exported from jax_sn.operators.transport_operator.

Submodules

• jax_sn.operators.transport_operator.implementations.multi_angle_inversions module
  • map_over_angles()
  • VmapAnglesForward
  • ScanAnglesForward
  • VmapAnglesInversion
  • ScanAnglesInversion
• jax_sn.operators.transport_operator.implementations.simplex_warp_kernel module
  • MAX_LENGTH
  • SUPPORTED
  • NDIM
  • INFLOW_EPSILON
  • SimplexWarpKernels
  • build_kernels()
• jax_sn.operators.transport_operator.implementations.transport_operator_forward module
  • VmapForward
  • BatchedForward
  • forward_transport_operator_vmap()
  • forward_transport_operator_batch()
• jax_sn.operators.transport_operator.implementations.transport_operator_jax module
  • VmapSweep
  • BatchedSweep
  • transport_inverse_sweep_batched()
  • transport_inverse_sweep_complete_vmap()
  • gij_kernel_idx_in_precomp()
  • gij_kernel_idx_in_precomp_vmap()
  • gij_inv_kernel_idx_in_precomp()
  • gij_inv_kernel_idx_in_precomp_vmap()
  • dij_f_p_fm_kernel_idx_in_precomp()
  • dij_f_p_fm_kernel_idx_in_precomp_vmap()
• jax_sn.operators.transport_operator.implementations.transport_operator_pallas module
  • transport_inverse_sweep_batched_pallas()
  • transport_inverse_sweep_complete_vmap_pallas()
  • sweep_pl_update()
  • transport_inverse_sweep_multiple_angles_batched_pallas()
  • VmapPallasSweep
  • BatchedPallasSweep
• jax_sn.operators.transport_operator.implementations.transport_operator_warp module
  • WarpSequentialSweep
  • WarpLevelSweep
  • WarpSequentialAngles
  • WarpLevelAngles