neuraLQX neuraLQX neuraLQX
/

Getting Started

  • Our Philosophy
  • Roadmap
  • Installing neuraLQX
  • Configuration Options
  • Citing neuraLQX
  • Quickstart

Contents

  • API Reference
    • neuralqx.graph package
      • neuralqx.graph.core package
        • neuralqx.graph.core.utils package
        • neuralqx.graph.core.abstract_graph module
        • neuralqx.graph.core.graph_handler module
      • neuralqx.graph.circular_ladder module
      • neuralqx.graph.dodecahedral module
      • neuralqx.graph.graph module
      • neuralqx.graph.grid_2d module
      • neuralqx.graph.half_ladder_graph module
      • neuralqx.graph.hexagonal_lattice module
      • neuralqx.graph.hypercube module
      • neuralqx.graph.k5_graph module
      • neuralqx.graph.ladder_graph module
      • neuralqx.graph.single_vertex module
      • neuralqx.graph.triangular_lattice module
    • neuralqx.hilbert package
      • neuralqx.hilbert.constraints package
        • neuralqx.hilbert.constraints.utils package
        • neuralqx.hilbert.constraints.u1_constraint module
      • neuralqx.hilbert.u1 package
        • neuralqx.hilbert.u1.index package
        • neuralqx.hilbert.u1.layout package
        • neuralqx.hilbert.u1.operations package
        • neuralqx.hilbert.u1.constrained_core module
        • neuralqx.hilbert.u1.interface module
        • neuralqx.hilbert.u1.unconstrained_core module
      • neuralqx.hilbert.utils package
        • neuralqx.hilbert.utils.index package
        • neuralqx.hilbert.utils.layout package
      • neuralqx.hilbert.abstract_hilbert_core module
      • neuralqx.hilbert.abstract_hilbert_interface module
    • neuralqx.operators package
      • neuralqx.operators.computational package
        • neuralqx.operators.computational.Euclidean3d package
        • neuralqx.operators.computational.Euclidean4d package
        • neuralqx.operators.computational.misc package
        • neuralqx.operators.computational.qr package
        • neuralqx.operators.computational.spherical package
        • neuralqx.operators.computational.wrappers package
      • neuralqx.operators.types package
        • neuralqx.operators.types.computational_operator package
        • neuralqx.operators.types.functional_local_operator package
      • neuralqx.operators.coloring module
      • neuralqx.operators.creation module
      • neuralqx.operators.holonomies module
      • neuralqx.operators.number module
    • neuralqx.gauge_groups package
      • neuralqx.gauge_groups.u1 package
        • neuralqx.gauge_groups.u1.u1_gauge_group module
      • neuralqx.gauge_groups.abstract_gauge_group_interface module
    • neuralqx.lqx package
      • neuralqx.lqx.qr package
        • neuralqx.lqx.qr.core package
        • neuralqx.lqx.qr.abstract_single_vertex_qr_model module
        • neuralqx.lqx.qr.single_vertex_qr_interface module
      • neuralqx.lqx.spherical package
        • neuralqx.lqx.spherical.core package
        • neuralqx.lqx.spherical.abstract_lqx_bosw_model module
        • neuralqx.lqx.spherical.spherical_bojowald_swiderski_interface module
      • neuralqx.lqx.wcl package
        • neuralqx.lqx.wcl.core package
        • neuralqx.lqx.wcl.abstract_lqx_wcl_model module
        • neuralqx.lqx.wcl.euclidean_interface module
      • neuralqx.lqx.abstract_lqx_interface module
      • neuralqx.lqx.abstract_lqx_model module
    • neuralqx.samplers package
      • neuralqx.samplers.rules package
      • neuralqx.samplers.builds module
      • neuralqx.samplers.export module
      • neuralqx.samplers.sampler module
      • neuralqx.samplers.types module
    • neuralqx.optimizer package
      • neuralqx.optimizer.solvers package
        • neuralqx.optimizer.solvers.builds module
        • neuralqx.optimizer.solvers.export module
        • neuralqx.optimizer.solvers.solvers module
        • neuralqx.optimizer.solvers.types module
      • neuralqx.optimizer.builds module
      • neuralqx.optimizer.export module
      • neuralqx.optimizer.optimizer module
      • neuralqx.optimizer.schedule module
      • neuralqx.optimizer.types module
    • neuralqx.nn package
      • neuralqx.nn.models package
        • neuralqx.nn.models.initialisers module
      • neuralqx.nn.projectors package
        • neuralqx.nn.projectors.group_projector package
        • neuralqx.nn.projectors.vacuum package
        • neuralqx.nn.projectors.utils module
    • neuralqx.solver package
      • neuralqx.solver.multi_state_solver module
      • neuralqx.solver.solver module
    • neuralqx.vqs package
      • neuralqx.vqs.mc package
        • neuralqx.vqs.mc.mc_state package
        • neuralqx.vqs.mc.common module
        • neuralqx.vqs.mc.kernels module
      • neuralqx.vqs.base module
    • neuralqx.callbacks package
      • neuralqx.callbacks.converged_state_export module
      • neuralqx.callbacks.live_monitoring module
      • neuralqx.callbacks.pairwise_gradient_cosine_callback module
    • neuralqx.utils package
      • neuralqx.utils.distributed package
        • neuralqx.utils.distributed.core module
        • neuralqx.utils.distributed.io module
      • neuralqx.utils.dtypes package
        • neuralqx.utils.dtypes.names module
        • neuralqx.utils.dtypes.policy module
        • neuralqx.utils.dtypes.runtime module
      • neuralqx.utils.io package
        • neuralqx.utils.io.loggers package
        • neuralqx.utils.io.printing package
        • neuralqx.utils.io.runtime_loggers package
      • neuralqx.utils.jax package
        • neuralqx.utils.jax.djaxrun module
        • neuralqx.utils.jax.sharding module
      • neuralqx.utils.misc package
        • neuralqx.utils.misc.arithmetic module
        • neuralqx.utils.misc.auth module
        • neuralqx.utils.misc.charges module
        • neuralqx.utils.misc.graph module
        • neuralqx.utils.misc.levi_civita module
        • neuralqx.utils.misc.modifiers module
        • neuralqx.utils.misc.triangulation module
      • neuralqx.utils.module package
        • neuralqx.utils.module.version package
        • neuralqx.utils.module.dependency_check module
        • neuralqx.utils.module.optional_dependency module
        • neuralqx.utils.module.version_check module
      • neuralqx.utils.symmetries package
        • neuralqx.utils.symmetries.types package
        • neuralqx.utils.symmetries.graph_symmetries module
        • neuralqx.utils.symmetries.utils module
      • neuralqx.utils.deprecation module
      • neuralqx.utils.errors module
      • neuralqx.utils.experimental module
      • neuralqx.utils.jit module
      • neuralqx.utils.numbers module
      • neuralqx.utils.parsing module
      • neuralqx.utils.serialization module
  • Tutorials
  • Guides
    • Graphs in neuraLQX
    • Hilbert spaces
    • Operators in neuraLQX
      • Matrix-based operators
      • Matrix-free Operators
      • Symbolic Operators
    • Gauge groups
    • LQX models
      • Euclidean weak-coupling models
      • Single-vertex QRLG
      • Spherical Bojowald–Swiderski model
    • Variational Quantum States
      • Penalty costs
      • Projectors and Projections
    • Solvers
    • Errors
    • Debugging in neuraLQX
    • Profiling in neuraLQX
  • Advanced Guides
    • Structs and data containers
    • Declarative Symbolic IR and Compiler Internals
      • Symbolic System Architecture
      • Symbolic IR Data Model
      • Compiler Pipeline and Caching
      • Extension Points and Customization
  • Parallelisation in neuraLQX
    • JAX Distributed Parallelisation (Single Host)
    • JAX Distributed Parallelisation (Multi-Host)
    • MPI (Deprecated)
      • CPU-based Parallelisation in neuraLQX
      • GPU-based Parallelisation in neuraLQX
      • Using neuraLQX in Parallel

Developers

  • Contributing to neuraLQX
  • Change Log
  • API Versioning and Stability
  • Experimental API
    • Single-trunk multi-head variational Monte Carlo (ST-MH)

On this page

  • replicate_sharding()
  • Submodules
  1. neuraLQX /
  2. API Reference /
  3. neuralqx.utils package /
  4. neuralqx.utils.jax package

neuralqx.utils.jax package¶

replicate_sharding(f_py)¶

Decorator replicating the sharding mechanism for jax.pure_callback-based get_conn_padded() implementations.

Submodules¶

  • neuralqx.utils.jax.djaxrun module
  • neuralqx.utils.jax.sharding module
Previous
neuralqx.utils.io.runtime_loggers.deferred_runtime_log module
Next
neuralqx.utils.jax.djaxrun module

2026, The neuraLQX Authors - All Rights Reserved

Made with Sphinx and Shibuya theme.