jax_utils

• jax_utils package
  • Submodules
  • jax_utils.common_tensor module
    • check_ndim_in()
    • is_broadcastable()
    • TensorAxes
      • TensorAxes.reverse_index()
      • TensorAxes.mandatory
      • TensorAxes.optional
    • expand_dims_axis()
    • AverageableArrayLike
      • AverageableArrayLike.mean()
    • Tensor
      • Tensor.array
      • Tensor.check_array()
      • Tensor.getitem_from_axes()
      • Tensor.values
      • Tensor.reverse_values()
      • Tensor.dtype
      • Tensor.shape
      • Tensor.ndim
      • Tensor.axes
      • Tensor.actual_axes
      • Tensor.mean()
      • Tensor.mean_over_axes()
      • Tensor.sum()
      • Tensor.sum_over_axes()
      • Tensor.has()
      • Tensor.has_actual()
      • Tensor.index()
      • Tensor.reverse_index()
      • Tensor.size()
      • Tensor.is_broadcastable_with()
    • RegularizedArrayLikeCost
      • RegularizedArrayLikeCost.cost
      • RegularizedArrayLikeCost.regularization
      • RegularizedArrayLikeCost.lagrangian_coefficient
      • RegularizedArrayLikeCost.mean()
  • jax_utils.compilation module
    • JaxCompilableProtocol
      • JaxCompilableProtocol.is_compilation_enabled
    • BaseJaxCompilable
      • BaseJaxCompilable.is_compilation_enabled
      • BaseJaxCompilable.enable_compilation()
      • BaseJaxCompilable.disable_compilation()
    • jit_when_compilation_enabled()
  • jax_utils.dynamics module
    • JaxDynamics
      • JaxDynamics.scalar_cost()
      • JaxDynamics.compute_gradient()
    • VectorizedJaxDynamics
      • VectorizedJaxDynamics.dynamics
      • VectorizedJaxDynamics.vectorized_axis
    • RegularizedJaxDynamics
      • RegularizedJaxDynamics.dynamics
      • RegularizedJaxDynamics.cost_regularizer
  • jax_utils.gradient module
    • BaseGradientStep
      • BaseGradientStep.init_optimizer()
      • BaseGradientStep.compute_gradient()
      • BaseGradientStep.update()
    • GradientStep
      • GradientStep.optimizer
      • GradientStep.dynamics
      • GradientStep.init_optimizer()
      • GradientStep.compute_gradient()
      • GradientStep.update()
    • VectorizedGradientStep
      • VectorizedGradientStep.gradient_step
      • VectorizedGradientStep.vectorized_axis
      • VectorizedGradientStep.init_optimizer()
      • VectorizedGradientStep.compute_gradient()
      • VectorizedGradientStep.update()
  • jax_utils.jax_tensor module
    • JaxTensor
      • JaxTensor.array
      • JaxTensor.is_broadcastable_with()
      • JaxTensor.from_flattened_axes()
      • JaxTensor.convert_to_axes()
    • NonNegativeValues
      • NonNegativeValues.array
      • NonNegativeValues.values
      • NonNegativeValues.reverse_values()
    • NonNegativeBudgetedValues
      • NonNegativeBudgetedValues.array
      • NonNegativeBudgetedValues.max_budget
      • NonNegativeBudgetedValues.values
      • NonNegativeBudgetedValues.reverse_values()
    • RegularizedJaxCost
      • RegularizedJaxCost.cost
      • RegularizedJaxCost.regularization
      • RegularizedJaxCost.lagrangian_coefficient
      • RegularizedJaxCost.mean()
  • jax_utils.markov_decision_process module
    • Dynamics
    • CostRegularizer
    • RegularizedDynamics
      • RegularizedDynamics.dynamics
      • RegularizedDynamics.cost_regularizer
  • jax_utils.optim module
    • OptimizationState
      • OptimizationState.iteration
      • OptimizationState.state
      • OptimizationState.action
      • OptimizationState.cost
      • OptimizationState.optimizer_state
      • OptimizationState.scalar_cost
    • OptimStoppingCondition
      • OptimStoppingCondition.stop()
      • OptimStoppingCondition.reset()
      • OptimStoppingCondition.nb_iterations_upper_bound
    • OptimStoppingConditionsCombination
      • OptimStoppingConditionsCombination.stopping_condition_1
      • OptimStoppingConditionsCombination.stopping_condition_2
      • OptimStoppingConditionsCombination.enable_compilation()
      • OptimStoppingConditionsCombination.disable_compilation()
    • OptimStoppingConditionIntersection
      • OptimStoppingConditionIntersection.stopping_condition_1
      • OptimStoppingConditionIntersection.stopping_condition_2
      • OptimStoppingConditionIntersection.stop()
      • OptimStoppingConditionIntersection.reset()
      • OptimStoppingConditionIntersection.nb_iterations_upper_bound
    • OptimStoppingConditionUnion
      • OptimStoppingConditionUnion.stopping_condition_1
      • OptimStoppingConditionUnion.stopping_condition_2
      • OptimStoppingConditionUnion.stop()
      • OptimStoppingConditionUnion.reset()
      • OptimStoppingConditionUnion.nb_iterations_upper_bound
    • MaxIterationsStoppingCondition
      • MaxIterationsStoppingCondition.max_iterations
      • MaxIterationsStoppingCondition.stop()
      • MaxIterationsStoppingCondition.nb_iterations_upper_bound
    • MinIterationsStoppingCondition
      • MinIterationsStoppingCondition.min_iterations
      • MinIterationsStoppingCondition.stop()
    • MinDeltaActionStoppingCondition
      • MinDeltaActionStoppingCondition.relative_tolerance
      • MinDeltaActionStoppingCondition.absolute_tolerance
      • MinDeltaActionStoppingCondition.stop()
      • MinDeltaActionStoppingCondition.reset()
    • MinDeltaCostStoppingCondition
      • MinDeltaCostStoppingCondition.relative_tolerance
      • MinDeltaCostStoppingCondition.absolute_tolerance
      • MinDeltaCostStoppingCondition.window_length
      • MinDeltaCostStoppingCondition.stop()
      • MinDeltaCostStoppingCondition.reset()
    • CostHistory
      • CostHistory.scalar_costs()
      • CostHistory.plot_scalar_costs()
    • GradientDescentOptimizationLoop
      • GradientDescentOptimizationLoop.gradient_step
      • GradientDescentOptimizationLoop.stopping_condition
      • GradientDescentOptimizationLoop.resume()
  • jax_utils.pytree module
    • ConvertibleToAxes
      • ConvertibleToAxes.convert_to_axes()
    • pytree_to_axes()
  • jax_utils.tranform module
    • scale_jax_tensor()
    • unscale_jax_tensor()
    • JaxScaler
      • JaxScaler.scaling_factors
      • JaxScaler.scaling_axes
      • JaxScaler.tensor_types_to_scale
      • JaxScaler.tensor_types_not_to_scale
      • JaxScaler.from_tensor()
      • JaxScaler.scale()
      • JaxScaler.unscale()
      • JaxScaler.expand_scaling_factors()
  • jax_utils.typing module
    • DataclassInstance
    • HashableIndexingOrSlicing
      • HashableIndexingOrSlicing.values
    • HashableSlicing
      • HashableSlicing.start
      • HashableSlicing.stop
      • HashableSlicing.step
      • HashableSlicing.values
    • HashableIndexing
      • HashableIndexing.indices
      • HashableIndexing.values
  • jax_utils.vectorization module
    • JaxVectorizableProtocol
      • JaxVectorizableProtocol.vectorized_axis
    • vectorize()
    • JaxDataclassNestedConvertibleToAxes
      • JaxDataclassNestedConvertibleToAxes.convert_to_axes()
  • Module contents