mirror of
https://github.com/DifferentiableUniverseInitiative/JaxPM.git
synced 2025-04-24 19:50:55 +00:00
34 lines
725 B
Python
34 lines
725 B
Python
import numpy as np
|
|
from _src.spmd_config import pm_operators
|
|
|
|
|
|
def fftn(arr):
|
|
return pm_operators.fftn(arr)
|
|
|
|
|
|
def ifftn(arr):
|
|
return pm_operators.ifftn(arr)
|
|
|
|
|
|
def halo_exchange(arr):
|
|
return pm_operators.halo_exchange(arr)
|
|
|
|
|
|
def slice_pad(arr, pad_width):
|
|
return pm_operators.slice_pad(arr, pad_width)
|
|
|
|
|
|
def slice_unpad(arr, pad_width):
|
|
return pm_operators.slice_unpad(arr, pad_width)
|
|
|
|
|
|
def normal(shape, key, dtype='float32'):
|
|
return pm_operators.normal(shape, key, dtype)
|
|
|
|
|
|
def fftk(shape, symmetric=True, finite=False, dtype=np.float32):
|
|
return pm_operators.fftk(shape, symmetric, finite, dtype)
|
|
|
|
|
|
def generate_initial_positions(shape):
|
|
return pm_operators.generate_initial_positions(shape)
|