From 052cacc3189c6fbe4d9a7837f5e9c210eedad1de Mon Sep 17 00:00:00 2001 From: Jeremy McGibbon Date: Tue, 6 Feb 2024 16:57:13 -0500 Subject: [PATCH] Revert "Rename package to ai2modulus" --- .pre-commit-config.yaml | 2 +- Dockerfile | 18 +++++++------- Makefile | 10 ++++---- docs/conf.py | 2 +- {ai2modulus => modulus}/__init__.py | 0 {ai2modulus => modulus}/constants.py | 0 {ai2modulus => modulus}/datapipes/__init__.py | 0 .../datapipes/benchmarks/__init__.py | 0 .../datapipes/benchmarks/darcy.py | 0 .../datapipes/benchmarks/kelvin_helmholtz.py | 0 .../datapipes/benchmarks/kernels/__init__.py | 0 .../benchmarks/kernels/finite_difference.py | 0 .../benchmarks/kernels/finite_volume.py | 0 .../datapipes/benchmarks/kernels/indexing.py | 0 .../benchmarks/kernels/initialization.py | 0 .../datapipes/benchmarks/kernels/utils.py | 0 .../datapipes/climate/__init__.py | 0 .../datapipes/climate/era5_hdf5.py | 0 .../datapipes/climate/era5_netcdf.py | 0 .../datapipes/climate/sfno/dataloader.py | 2 +- .../sfno/dataloaders/dali_es_helper_2d.py | 2 +- .../sfno/dataloaders/data_loader_dali_2d.py | 4 ++-- .../sfno/dataloaders/data_loader_dummy.py | 2 +- .../dataloaders/data_loader_multifiles.py | 2 +- .../climate/sfno/dataloaders/zarr_helper.py | 0 {ai2modulus => modulus}/datapipes/datapipe.py | 2 +- .../datapipes/gnn/mgn_dataset.py | 0 {ai2modulus => modulus}/datapipes/meta.py | 0 .../deploy/onnx/__init__.py | 0 {ai2modulus => modulus}/deploy/onnx/utils.py | 0 .../deploy/triton/__init__.py | 0 .../deploy/trt/__init__.py | 0 .../distributed/__init__.py | 0 .../distributed/manager.py | 0 {ai2modulus => modulus}/distributed/utils.py | 0 {ai2modulus => modulus}/metrics/__init__.py | 0 .../metrics/climate/__init__.py | 0 .../metrics/climate/acc.py | 2 +- .../metrics/climate/efi.py | 4 ++-- .../metrics/climate/reduction.py | 2 +- .../metrics/general/__init__.py | 0 .../metrics/general/calibration.py | 2 +- .../metrics/general/crps.py | 0 .../metrics/general/ensemble_metrics.py | 2 +- .../metrics/general/entropy.py | 0 .../metrics/general/histogram.py | 2 +- .../metrics/general/mse.py | 0 .../metrics/general/reduction.py | 0 .../metrics/general/wasserstein.py | 0 {ai2modulus => modulus}/models/__init__.py | 0 .../models/afno/__init__.py | 0 {ai2modulus => modulus}/models/afno/afno.py | 4 ++-- .../models/afno/distributed/__init__.py | 0 .../models/afno/distributed/afno.py | 0 .../models/afno/distributed/layers.py | 0 .../models/afno/distributed/mappings.py | 0 .../models/fno/__init__.py | 0 {ai2modulus => modulus}/models/fno/fno.py | 9 +++---- .../models/gnn_layers/__init__.py | 0 .../models/gnn_layers/embedder.py | 0 .../models/gnn_layers/mesh_edge_block.py | 0 .../models/gnn_layers/mesh_graph_decoder.py | 0 .../models/gnn_layers/mesh_graph_encoder.py | 0 .../models/gnn_layers/mesh_graph_mlp.py | 2 +- .../models/gnn_layers/mesh_node_block.py | 0 .../models/gnn_layers/utils.py | 0 .../models/graphcast/__init__.py | 0 .../models/graphcast/graph_cast_net.py | 18 +++++++------- .../models/graphcast/graph_cast_processor.py | 6 ++--- .../models/layers/__init__.py | 0 .../models/layers/activations.py | 2 +- {ai2modulus => modulus}/models/layers/fft.py | 0 .../models/layers/fully_connected_layers.py | 0 .../models/layers/fused_silu.py | 0 .../models/layers/spectral_layers.py | 0 .../models/layers/weight_norm.py | 2 +- .../models/meshgraphnet/__init__.py | 0 .../models/meshgraphnet/meshgraphnet.py | 14 +++++------ {ai2modulus => modulus}/models/meta.py | 0 .../models/mlp/__init__.py | 0 .../models/mlp/fully_connected.py | 4 ++-- {ai2modulus => modulus}/models/module.py | 2 +- .../models/pix2pix/__init__.py | 0 .../models/pix2pix/pix2pix.py | 2 +- .../models/rnn/__init__.py | 0 {ai2modulus => modulus}/models/rnn/layers.py | 0 .../models/rnn/rnn_one2many.py | 8 +++---- .../models/rnn/rnn_seq2seq.py | 8 +++---- .../models/sfno/activations.py | 0 .../models/sfno/contractions.py | 0 .../models/sfno/factorizations.py | 2 +- .../models/sfno/initialization.py | 0 {ai2modulus => modulus}/models/sfno/layers.py | 6 ++--- .../models/sfno/preprocessor.py | 0 .../models/sfno/s2convolutions.py | 10 ++++---- .../models/sfno/sfnonet.py | 18 +++++++------- .../models/srrn/__init__.py | 0 .../models/srrn/super_res_net.py | 2 +- {ai2modulus => modulus}/utils/__init__.py | 0 {ai2modulus => modulus}/utils/capture.py | 24 +++++++++---------- .../utils/graphcast/__init__.py | 0 .../utils/graphcast/data_utils.py | 0 .../utils/graphcast/graph.py | 0 .../utils/graphcast/graph_utils.py | 0 .../utils/graphcast/icospheres.py | 0 .../utils/graphcast/loss.py | 0 {ai2modulus => modulus}/utils/sfno/YParams.py | 0 .../utils/sfno/distributed/__init__.py | 0 .../utils/sfno/distributed/comm.py | 2 +- .../utils/sfno/distributed/helpers.py | 2 +- .../utils/sfno/distributed/layer_norm.py | 4 ++-- .../utils/sfno/distributed/layers.py | 18 +++++++------- .../utils/sfno/distributed/mappings.py | 8 +++---- .../utils/sfno/img_utils.py | 0 .../utils/sfno/logging_utils.py | 0 {ai2modulus => modulus}/utils/sfno/loss.py | 4 ++-- {ai2modulus => modulus}/utils/sfno/metric.py | 6 ++--- .../utils/sfno/metrics/weighted_acc_rmse.py | 6 ++--- .../sfno/metrics/weighted_acc_rmse_inf.py | 2 +- .../utils/sfno/warmup_scheduler.py | 0 .../utils/sfno/zenith_angle.py | 0 pyproject.toml | 6 ++--- test/datapipes/common/cuda_graphs.py | 8 +++---- test/datapipes/common/device.py | 2 ++ test/datapipes/common/iterator_check.py | 4 ++-- test/datapipes/common/sequence.py | 2 ++ test/datapipes/common/shape_check.py | 1 + test/datapipes/common/shuffle.py | 2 ++ test/datapipes/sfno/test_data_loader_dummy.py | 2 +- test/datapipes/test_darcy.py | 2 +- test/datapipes/test_era5_hdf5.py | 2 +- test/datapipes/test_kelvin_helmholtz.py | 2 +- test/deploy/test_onnx_fft.py | 4 ++-- test/deploy/test_onnx_utils.py | 4 ++-- test/distributed/test_manager.py | 2 +- test/distributed/test_utils.py | 4 ++-- test/metrics/test_metrics_climate.py | 12 +++++----- test/metrics/test_metrics_general.py | 16 ++++++------- test/models/common/checkpoints.py | 14 +++++------ test/models/common/fwdaccuracy.py | 6 ++--- test/models/common/inference.py | 12 +++++----- test/models/common/optimization.py | 18 +++++++------- test/models/graphcast/test_concat_trick.py | 2 +- test/models/graphcast/test_cugraphops.py | 2 +- .../graphcast/test_grad_checkpointing.py | 2 +- test/models/graphcast/test_graphcast.py | 2 +- test/models/sfno/test_activations.py | 2 +- test/models/sfno/test_sfno.py | 2 +- test/models/test_afno.py | 2 +- test/models/test_fno.py | 4 ++-- test/models/test_fully_connected.py | 2 +- test/models/test_layers_activations.py | 4 ++-- test/models/test_layers_weightnorm.py | 2 +- test/models/test_meshgraphnet.py | 2 +- test/models/test_pix2pix.py | 2 +- test/models/test_rnn.py | 4 ++-- test/models/test_rnn_layers.py | 2 +- test/models/test_super_res_net.py | 2 +- .../graphcast/test_coordinate_transform.py | 2 +- test/utils/graphcast/test_loss.py | 2 +- test/utils/sfno/test_img_utils.py | 2 +- test/utils/sfno/test_logging.py | 2 +- test/utils/sfno/test_warmup_scheduler.py | 2 +- test/utils/sfno/test_yparams.py | 2 +- test/utils/test_capture.py | 6 ++--- 165 files changed, 222 insertions(+), 214 deletions(-) rename {ai2modulus => modulus}/__init__.py (100%) rename {ai2modulus => modulus}/constants.py (100%) rename {ai2modulus => modulus}/datapipes/__init__.py (100%) rename {ai2modulus => modulus}/datapipes/benchmarks/__init__.py (100%) rename {ai2modulus => modulus}/datapipes/benchmarks/darcy.py (100%) rename {ai2modulus => modulus}/datapipes/benchmarks/kelvin_helmholtz.py (100%) rename {ai2modulus => modulus}/datapipes/benchmarks/kernels/__init__.py (100%) rename {ai2modulus => modulus}/datapipes/benchmarks/kernels/finite_difference.py (100%) rename {ai2modulus => modulus}/datapipes/benchmarks/kernels/finite_volume.py (100%) rename {ai2modulus => modulus}/datapipes/benchmarks/kernels/indexing.py (100%) rename {ai2modulus => modulus}/datapipes/benchmarks/kernels/initialization.py (100%) rename {ai2modulus => modulus}/datapipes/benchmarks/kernels/utils.py (100%) rename {ai2modulus => modulus}/datapipes/climate/__init__.py (100%) rename {ai2modulus => modulus}/datapipes/climate/era5_hdf5.py (100%) rename {ai2modulus => modulus}/datapipes/climate/era5_netcdf.py (100%) rename {ai2modulus => modulus}/datapipes/climate/sfno/dataloader.py (99%) rename {ai2modulus => modulus}/datapipes/climate/sfno/dataloaders/dali_es_helper_2d.py (99%) rename {ai2modulus => modulus}/datapipes/climate/sfno/dataloaders/data_loader_dali_2d.py (99%) rename {ai2modulus => modulus}/datapipes/climate/sfno/dataloaders/data_loader_dummy.py (99%) rename {ai2modulus => modulus}/datapipes/climate/sfno/dataloaders/data_loader_multifiles.py (98%) rename {ai2modulus => modulus}/datapipes/climate/sfno/dataloaders/zarr_helper.py (100%) rename {ai2modulus => modulus}/datapipes/datapipe.py (97%) rename {ai2modulus => modulus}/datapipes/gnn/mgn_dataset.py (100%) rename {ai2modulus => modulus}/datapipes/meta.py (100%) rename {ai2modulus => modulus}/deploy/onnx/__init__.py (100%) rename {ai2modulus => modulus}/deploy/onnx/utils.py (100%) rename {ai2modulus => modulus}/deploy/triton/__init__.py (100%) rename {ai2modulus => modulus}/deploy/trt/__init__.py (100%) rename {ai2modulus => modulus}/distributed/__init__.py (100%) rename {ai2modulus => modulus}/distributed/manager.py (100%) rename {ai2modulus => modulus}/distributed/utils.py (100%) rename {ai2modulus => modulus}/metrics/__init__.py (100%) rename {ai2modulus => modulus}/metrics/climate/__init__.py (100%) rename {ai2modulus => modulus}/metrics/climate/acc.py (97%) rename {ai2modulus => modulus}/metrics/climate/efi.py (95%) rename {ai2modulus => modulus}/metrics/climate/reduction.py (98%) rename {ai2modulus => modulus}/metrics/general/__init__.py (100%) rename {ai2modulus => modulus}/metrics/general/calibration.py (98%) rename {ai2modulus => modulus}/metrics/general/crps.py (100%) rename {ai2modulus => modulus}/metrics/general/ensemble_metrics.py (99%) rename {ai2modulus => modulus}/metrics/general/entropy.py (100%) rename {ai2modulus => modulus}/metrics/general/histogram.py (99%) rename {ai2modulus => modulus}/metrics/general/mse.py (100%) rename {ai2modulus => modulus}/metrics/general/reduction.py (100%) rename {ai2modulus => modulus}/metrics/general/wasserstein.py (100%) rename {ai2modulus => modulus}/models/__init__.py (100%) rename {ai2modulus => modulus}/models/afno/__init__.py (100%) rename {ai2modulus => modulus}/models/afno/afno.py (99%) rename {ai2modulus => modulus}/models/afno/distributed/__init__.py (100%) rename {ai2modulus => modulus}/models/afno/distributed/afno.py (100%) rename {ai2modulus => modulus}/models/afno/distributed/layers.py (100%) rename {ai2modulus => modulus}/models/afno/distributed/mappings.py (100%) rename {ai2modulus => modulus}/models/fno/__init__.py (100%) rename {ai2modulus => modulus}/models/fno/fno.py (99%) rename {ai2modulus => modulus}/models/gnn_layers/__init__.py (100%) rename {ai2modulus => modulus}/models/gnn_layers/embedder.py (100%) rename {ai2modulus => modulus}/models/gnn_layers/mesh_edge_block.py (100%) rename {ai2modulus => modulus}/models/gnn_layers/mesh_graph_decoder.py (100%) rename {ai2modulus => modulus}/models/gnn_layers/mesh_graph_encoder.py (100%) rename {ai2modulus => modulus}/models/gnn_layers/mesh_graph_mlp.py (99%) rename {ai2modulus => modulus}/models/gnn_layers/mesh_node_block.py (100%) rename {ai2modulus => modulus}/models/gnn_layers/utils.py (100%) rename {ai2modulus => modulus}/models/graphcast/__init__.py (100%) rename {ai2modulus => modulus}/models/graphcast/graph_cast_net.py (97%) rename {ai2modulus => modulus}/models/graphcast/graph_cast_processor.py (96%) rename {ai2modulus => modulus}/models/layers/__init__.py (100%) rename {ai2modulus => modulus}/models/layers/activations.py (99%) rename {ai2modulus => modulus}/models/layers/fft.py (100%) rename {ai2modulus => modulus}/models/layers/fully_connected_layers.py (100%) rename {ai2modulus => modulus}/models/layers/fused_silu.py (100%) rename {ai2modulus => modulus}/models/layers/spectral_layers.py (100%) rename {ai2modulus => modulus}/models/layers/weight_norm.py (99%) rename {ai2modulus => modulus}/models/meshgraphnet/__init__.py (100%) rename {ai2modulus => modulus}/models/meshgraphnet/meshgraphnet.py (95%) rename {ai2modulus => modulus}/models/meta.py (100%) rename {ai2modulus => modulus}/models/mlp/__init__.py (100%) rename {ai2modulus => modulus}/models/mlp/fully_connected.py (98%) rename {ai2modulus => modulus}/models/module.py (98%) rename {ai2modulus => modulus}/models/pix2pix/__init__.py (100%) rename {ai2modulus => modulus}/models/pix2pix/pix2pix.py (99%) rename {ai2modulus => modulus}/models/rnn/__init__.py (100%) rename {ai2modulus => modulus}/models/rnn/layers.py (100%) rename {ai2modulus => modulus}/models/rnn/rnn_one2many.py (98%) rename {ai2modulus => modulus}/models/rnn/rnn_seq2seq.py (98%) rename {ai2modulus => modulus}/models/sfno/activations.py (100%) rename {ai2modulus => modulus}/models/sfno/contractions.py (100%) rename {ai2modulus => modulus}/models/sfno/factorizations.py (99%) rename {ai2modulus => modulus}/models/sfno/initialization.py (100%) rename {ai2modulus => modulus}/models/sfno/layers.py (99%) rename {ai2modulus => modulus}/models/sfno/preprocessor.py (100%) rename {ai2modulus => modulus}/models/sfno/s2convolutions.py (98%) rename {ai2modulus => modulus}/models/sfno/sfnonet.py (98%) rename {ai2modulus => modulus}/models/srrn/__init__.py (100%) rename {ai2modulus => modulus}/models/srrn/super_res_net.py (99%) rename {ai2modulus => modulus}/utils/__init__.py (100%) rename {ai2modulus => modulus}/utils/capture.py (95%) rename {ai2modulus => modulus}/utils/graphcast/__init__.py (100%) rename {ai2modulus => modulus}/utils/graphcast/data_utils.py (100%) rename {ai2modulus => modulus}/utils/graphcast/graph.py (100%) rename {ai2modulus => modulus}/utils/graphcast/graph_utils.py (100%) rename {ai2modulus => modulus}/utils/graphcast/icospheres.py (100%) rename {ai2modulus => modulus}/utils/graphcast/loss.py (100%) rename {ai2modulus => modulus}/utils/sfno/YParams.py (100%) rename {ai2modulus => modulus}/utils/sfno/distributed/__init__.py (100%) rename {ai2modulus => modulus}/utils/sfno/distributed/comm.py (99%) rename {ai2modulus => modulus}/utils/sfno/distributed/helpers.py (99%) rename {ai2modulus => modulus}/utils/sfno/distributed/layer_norm.py (97%) rename {ai2modulus => modulus}/utils/sfno/distributed/layers.py (96%) rename {ai2modulus => modulus}/utils/sfno/distributed/mappings.py (97%) rename {ai2modulus => modulus}/utils/sfno/img_utils.py (100%) rename {ai2modulus => modulus}/utils/sfno/logging_utils.py (100%) rename {ai2modulus => modulus}/utils/sfno/loss.py (99%) rename {ai2modulus => modulus}/utils/sfno/metric.py (98%) rename {ai2modulus => modulus}/utils/sfno/metrics/weighted_acc_rmse.py (97%) rename {ai2modulus => modulus}/utils/sfno/metrics/weighted_acc_rmse_inf.py (99%) rename {ai2modulus => modulus}/utils/sfno/warmup_scheduler.py (100%) rename {ai2modulus => modulus}/utils/sfno/zenith_angle.py (100%) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index dd6e7b4761..760d025a99 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -29,7 +29,7 @@ repos: "--ignore-magic", "--fail-under=99", "--exclude=['setup.py', 'test', 'build', 'docs']", "--ignore-regex=['forward', 'reset_parameters', 'extra_repr', 'MetaData', 'apply_activation','exec_activation']", "--color", "--"] - exclude: ^ai2modulus/internal/ + exclude: ^modulus/internal/ - repo: local hooks: diff --git a/Dockerfile b/Dockerfile index ef58e62902..e3d8c150d4 100644 --- a/Dockerfile +++ b/Dockerfile @@ -28,16 +28,16 @@ ENV _CUDA_COMPAT_TIMEOUT=90 # Install custom onnx # TODO: Find a fix to eliminate the custom build # Forcing numpy update to over ride numba 0.56.4 max numpy constraint -COPY . /ai2modulus/ -RUN if [ -e "/ai2modulus/deps/onnxruntime_gpu-1.14.0-cp38-cp38-linux_x86_64.whl" ]; then \ +COPY . /modulus/ +RUN if [ -e "/modulus/deps/onnxruntime_gpu-1.14.0-cp38-cp38-linux_x86_64.whl" ]; then \ echo "Custom wheel exists, installing!" && \ - pip install --force-reinstall /ai2modulus/deps/onnxruntime_gpu-1.14.0-cp38-cp38-linux_x86_64.whl; \ + pip install --force-reinstall /modulus/deps/onnxruntime_gpu-1.14.0-cp38-cp38-linux_x86_64.whl; \ else \ echo "No custom wheel present, skipping" && \ pip install numpy==1.22.4; \ fi # cleanup of stage -RUN rm -rf /ai2modulus/ +RUN rm -rf /modulus/ # CI image FROM builder as ci @@ -63,17 +63,17 @@ RUN mkdir -p /opt/cugraphops &&\ ENV PYTHONPATH="${PYTHONPATH}:/opt/cugraphops/lib/python3.8/site-packages" -COPY . /ai2modulus/ -RUN cd /ai2modulus/ && pip install -e . && rm -rf /ai2modulus/ +COPY . /modulus/ +RUN cd /modulus/ && pip install -e . && rm -rf /modulus/ # Deployment image FROM builder as deploy RUN pip install protobuf==3.20.0 -COPY . /ai2modulus/ -RUN cd /ai2modulus/ && pip install . +COPY . /modulus/ +RUN cd /modulus/ && pip install . # Clean up -RUN rm -rf /ai2modulus/ +RUN rm -rf /modulus/ # Docs image FROM deploy as docs diff --git a/Makefile b/Makefile index c7438b5f53..d1e00f66bd 100644 --- a/Makefile +++ b/Makefile @@ -12,7 +12,7 @@ black: black --check --exclude=docs/ ./ interrogate: - cd ai2modulus && \ + cd modulus && \ interrogate --ignore-init-method \ --ignore-init-module \ --ignore-module \ @@ -41,7 +41,7 @@ doctest: coverage run \ --rcfile='test/coverage.docstring.rc' \ -m pytest \ - --doctest-modules ai2modulus/ --ignore-glob=*internal* + --doctest-modules modulus/ --ignore-glob=*internal* pytest: coverage run \ @@ -59,11 +59,11 @@ coverage: coverage html container-deploy: - docker build -t ai2modulus:deploy --target deploy -f Dockerfile . + docker build -t modulus:deploy --target deploy -f Dockerfile . container-ci: - docker build -t ai2modulus:ci --target ci -f Dockerfile . + docker build -t modulus:ci --target ci -f Dockerfile . container-docs: - docker build -t ai2modulus:docs --target docs -f Dockerfile . + docker build -t modulus:docs --target docs -f Dockerfile . diff --git a/docs/conf.py b/docs/conf.py index b27a5c39d3..03d961edf4 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -21,7 +21,7 @@ # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information import os import sphinx_rtd_theme -from ai2modulus import __version__ as version +from modulus import __version__ as version project = "NVIDIA Modulus" copyright = "2023, NVIDIA Modulus Team" diff --git a/ai2modulus/__init__.py b/modulus/__init__.py similarity index 100% rename from ai2modulus/__init__.py rename to modulus/__init__.py diff --git a/ai2modulus/constants.py b/modulus/constants.py similarity index 100% rename from ai2modulus/constants.py rename to modulus/constants.py diff --git a/ai2modulus/datapipes/__init__.py b/modulus/datapipes/__init__.py similarity index 100% rename from ai2modulus/datapipes/__init__.py rename to modulus/datapipes/__init__.py diff --git a/ai2modulus/datapipes/benchmarks/__init__.py b/modulus/datapipes/benchmarks/__init__.py similarity index 100% rename from ai2modulus/datapipes/benchmarks/__init__.py rename to modulus/datapipes/benchmarks/__init__.py diff --git a/ai2modulus/datapipes/benchmarks/darcy.py b/modulus/datapipes/benchmarks/darcy.py similarity index 100% rename from ai2modulus/datapipes/benchmarks/darcy.py rename to modulus/datapipes/benchmarks/darcy.py diff --git a/ai2modulus/datapipes/benchmarks/kelvin_helmholtz.py b/modulus/datapipes/benchmarks/kelvin_helmholtz.py similarity index 100% rename from ai2modulus/datapipes/benchmarks/kelvin_helmholtz.py rename to modulus/datapipes/benchmarks/kelvin_helmholtz.py diff --git a/ai2modulus/datapipes/benchmarks/kernels/__init__.py b/modulus/datapipes/benchmarks/kernels/__init__.py similarity index 100% rename from ai2modulus/datapipes/benchmarks/kernels/__init__.py rename to modulus/datapipes/benchmarks/kernels/__init__.py diff --git a/ai2modulus/datapipes/benchmarks/kernels/finite_difference.py b/modulus/datapipes/benchmarks/kernels/finite_difference.py similarity index 100% rename from ai2modulus/datapipes/benchmarks/kernels/finite_difference.py rename to modulus/datapipes/benchmarks/kernels/finite_difference.py diff --git a/ai2modulus/datapipes/benchmarks/kernels/finite_volume.py b/modulus/datapipes/benchmarks/kernels/finite_volume.py similarity index 100% rename from ai2modulus/datapipes/benchmarks/kernels/finite_volume.py rename to modulus/datapipes/benchmarks/kernels/finite_volume.py diff --git a/ai2modulus/datapipes/benchmarks/kernels/indexing.py b/modulus/datapipes/benchmarks/kernels/indexing.py similarity index 100% rename from ai2modulus/datapipes/benchmarks/kernels/indexing.py rename to modulus/datapipes/benchmarks/kernels/indexing.py diff --git a/ai2modulus/datapipes/benchmarks/kernels/initialization.py b/modulus/datapipes/benchmarks/kernels/initialization.py similarity index 100% rename from ai2modulus/datapipes/benchmarks/kernels/initialization.py rename to modulus/datapipes/benchmarks/kernels/initialization.py diff --git a/ai2modulus/datapipes/benchmarks/kernels/utils.py b/modulus/datapipes/benchmarks/kernels/utils.py similarity index 100% rename from ai2modulus/datapipes/benchmarks/kernels/utils.py rename to modulus/datapipes/benchmarks/kernels/utils.py diff --git a/ai2modulus/datapipes/climate/__init__.py b/modulus/datapipes/climate/__init__.py similarity index 100% rename from ai2modulus/datapipes/climate/__init__.py rename to modulus/datapipes/climate/__init__.py diff --git a/ai2modulus/datapipes/climate/era5_hdf5.py b/modulus/datapipes/climate/era5_hdf5.py similarity index 100% rename from ai2modulus/datapipes/climate/era5_hdf5.py rename to modulus/datapipes/climate/era5_hdf5.py diff --git a/ai2modulus/datapipes/climate/era5_netcdf.py b/modulus/datapipes/climate/era5_netcdf.py similarity index 100% rename from ai2modulus/datapipes/climate/era5_netcdf.py rename to modulus/datapipes/climate/era5_netcdf.py diff --git a/ai2modulus/datapipes/climate/sfno/dataloader.py b/modulus/datapipes/climate/sfno/dataloader.py similarity index 99% rename from ai2modulus/datapipes/climate/sfno/dataloader.py rename to modulus/datapipes/climate/sfno/dataloader.py index 58f4a17e89..839a0618db 100644 --- a/ai2modulus/datapipes/climate/sfno/dataloader.py +++ b/modulus/datapipes/climate/sfno/dataloader.py @@ -20,7 +20,7 @@ # distributed stuff import torch.distributed as dist -from ai2modulus.utils.sfno.distributed import comm +from modulus.utils.sfno.distributed import comm def init_distributed_io(params): # pragma: no cover diff --git a/ai2modulus/datapipes/climate/sfno/dataloaders/dali_es_helper_2d.py b/modulus/datapipes/climate/sfno/dataloaders/dali_es_helper_2d.py similarity index 99% rename from ai2modulus/datapipes/climate/sfno/dataloaders/dali_es_helper_2d.py rename to modulus/datapipes/climate/sfno/dataloaders/dali_es_helper_2d.py index 21d0f8d64e..a09f0b3a3a 100644 --- a/ai2modulus/datapipes/climate/sfno/dataloaders/dali_es_helper_2d.py +++ b/modulus/datapipes/climate/sfno/dataloaders/dali_es_helper_2d.py @@ -29,7 +29,7 @@ # we need this for the zenith angle feature import datetime -from ai2modulus.utils.sfno.zenith_angle import cos_zenith_angle +from modulus.utils.sfno.zenith_angle import cos_zenith_angle class GeneralES(object): diff --git a/ai2modulus/datapipes/climate/sfno/dataloaders/data_loader_dali_2d.py b/modulus/datapipes/climate/sfno/dataloaders/data_loader_dali_2d.py similarity index 99% rename from ai2modulus/datapipes/climate/sfno/dataloaders/data_loader_dali_2d.py rename to modulus/datapipes/climate/sfno/dataloaders/data_loader_dali_2d.py index 0482f694e1..dea2ab409f 100644 --- a/ai2modulus/datapipes/climate/sfno/dataloaders/data_loader_dali_2d.py +++ b/modulus/datapipes/climate/sfno/dataloaders/data_loader_dali_2d.py @@ -16,7 +16,7 @@ import numpy as np # distributed stuff -from ai2modulus.utils.sfno.distributed import comm +from modulus.utils.sfno.distributed import comm # DALI stuff from nvidia.dali.pipeline import Pipeline @@ -25,7 +25,7 @@ from nvidia.dali.plugin.pytorch import DALIGenericIterator, LastBatchPolicy # es helper -import ai2modulus.datapipes.climate.sfno.dataloaders.dali_es_helper_2d as esh +import modulus.datapipes.climate.sfno.dataloaders.dali_es_helper_2d as esh class ERA5DaliESDataloader(object): diff --git a/ai2modulus/datapipes/climate/sfno/dataloaders/data_loader_dummy.py b/modulus/datapipes/climate/sfno/dataloaders/data_loader_dummy.py similarity index 99% rename from ai2modulus/datapipes/climate/sfno/dataloaders/data_loader_dummy.py rename to modulus/datapipes/climate/sfno/dataloaders/data_loader_dummy.py index 1f0ed26f96..4e87e4efbe 100644 --- a/ai2modulus/datapipes/climate/sfno/dataloaders/data_loader_dummy.py +++ b/modulus/datapipes/climate/sfno/dataloaders/data_loader_dummy.py @@ -19,7 +19,7 @@ import h5py # distributed stuff -from ai2modulus.utils.sfno.distributed import comm +from modulus.utils.sfno.distributed import comm class DummyLoader(object): diff --git a/ai2modulus/datapipes/climate/sfno/dataloaders/data_loader_multifiles.py b/modulus/datapipes/climate/sfno/dataloaders/data_loader_multifiles.py similarity index 98% rename from ai2modulus/datapipes/climate/sfno/dataloaders/data_loader_multifiles.py rename to modulus/datapipes/climate/sfno/dataloaders/data_loader_multifiles.py index 80e2e0e76e..c3113a382a 100644 --- a/ai2modulus/datapipes/climate/sfno/dataloaders/data_loader_multifiles.py +++ b/modulus/datapipes/climate/sfno/dataloaders/data_loader_multifiles.py @@ -20,7 +20,7 @@ import h5py # import cv2 -from ai2modulus.utils.sfno.img_utils import reshape_fields +from modulus.utils.sfno.img_utils import reshape_fields class MultifilesDataset(Dataset): diff --git a/ai2modulus/datapipes/climate/sfno/dataloaders/zarr_helper.py b/modulus/datapipes/climate/sfno/dataloaders/zarr_helper.py similarity index 100% rename from ai2modulus/datapipes/climate/sfno/dataloaders/zarr_helper.py rename to modulus/datapipes/climate/sfno/dataloaders/zarr_helper.py diff --git a/ai2modulus/datapipes/datapipe.py b/modulus/datapipes/datapipe.py similarity index 97% rename from ai2modulus/datapipes/datapipe.py rename to modulus/datapipes/datapipe.py index 7048e56335..4e40c41cd3 100644 --- a/ai2modulus/datapipes/datapipe.py +++ b/modulus/datapipes/datapipe.py @@ -17,7 +17,7 @@ from typing import Union from pathlib import Path -from ai2modulus.datapipes.meta import DatapipeMetaData +from modulus.datapipes.meta import DatapipeMetaData class Datapipe: diff --git a/ai2modulus/datapipes/gnn/mgn_dataset.py b/modulus/datapipes/gnn/mgn_dataset.py similarity index 100% rename from ai2modulus/datapipes/gnn/mgn_dataset.py rename to modulus/datapipes/gnn/mgn_dataset.py diff --git a/ai2modulus/datapipes/meta.py b/modulus/datapipes/meta.py similarity index 100% rename from ai2modulus/datapipes/meta.py rename to modulus/datapipes/meta.py diff --git a/ai2modulus/deploy/onnx/__init__.py b/modulus/deploy/onnx/__init__.py similarity index 100% rename from ai2modulus/deploy/onnx/__init__.py rename to modulus/deploy/onnx/__init__.py diff --git a/ai2modulus/deploy/onnx/utils.py b/modulus/deploy/onnx/utils.py similarity index 100% rename from ai2modulus/deploy/onnx/utils.py rename to modulus/deploy/onnx/utils.py diff --git a/ai2modulus/deploy/triton/__init__.py b/modulus/deploy/triton/__init__.py similarity index 100% rename from ai2modulus/deploy/triton/__init__.py rename to modulus/deploy/triton/__init__.py diff --git a/ai2modulus/deploy/trt/__init__.py b/modulus/deploy/trt/__init__.py similarity index 100% rename from ai2modulus/deploy/trt/__init__.py rename to modulus/deploy/trt/__init__.py diff --git a/ai2modulus/distributed/__init__.py b/modulus/distributed/__init__.py similarity index 100% rename from ai2modulus/distributed/__init__.py rename to modulus/distributed/__init__.py diff --git a/ai2modulus/distributed/manager.py b/modulus/distributed/manager.py similarity index 100% rename from ai2modulus/distributed/manager.py rename to modulus/distributed/manager.py diff --git a/ai2modulus/distributed/utils.py b/modulus/distributed/utils.py similarity index 100% rename from ai2modulus/distributed/utils.py rename to modulus/distributed/utils.py diff --git a/ai2modulus/metrics/__init__.py b/modulus/metrics/__init__.py similarity index 100% rename from ai2modulus/metrics/__init__.py rename to modulus/metrics/__init__.py diff --git a/ai2modulus/metrics/climate/__init__.py b/modulus/metrics/climate/__init__.py similarity index 100% rename from ai2modulus/metrics/climate/__init__.py rename to modulus/metrics/climate/__init__.py diff --git a/ai2modulus/metrics/climate/acc.py b/modulus/metrics/climate/acc.py similarity index 97% rename from ai2modulus/metrics/climate/acc.py rename to modulus/metrics/climate/acc.py index d92c23a282..d3cf470ef5 100644 --- a/ai2modulus/metrics/climate/acc.py +++ b/modulus/metrics/climate/acc.py @@ -15,7 +15,7 @@ # TODO(Dallas) Introduce Distributed Class for computation. import torch -from ai2modulus.metrics.climate.reduction import _compute_lat_weights +from modulus.metrics.climate.reduction import _compute_lat_weights Tensor = torch.Tensor diff --git a/ai2modulus/metrics/climate/efi.py b/modulus/metrics/climate/efi.py similarity index 95% rename from ai2modulus/metrics/climate/efi.py rename to modulus/metrics/climate/efi.py index f9ca43ea50..64884cc719 100644 --- a/ai2modulus/metrics/climate/efi.py +++ b/modulus/metrics/climate/efi.py @@ -15,8 +15,8 @@ # TODO(Dallas) Introduce Distributed Class for computation. import torch -from ai2modulus.metrics.general.histogram import normal_pdf, normal_cdf, histogram -from ai2modulus.metrics.general.entropy import _entropy_from_counts +from modulus.metrics.general.histogram import normal_pdf, normal_cdf, histogram +from modulus.metrics.general.entropy import _entropy_from_counts Tensor = torch.Tensor diff --git a/ai2modulus/metrics/climate/reduction.py b/modulus/metrics/climate/reduction.py similarity index 98% rename from ai2modulus/metrics/climate/reduction.py rename to modulus/metrics/climate/reduction.py index 0cdb07b79b..3cdb7fbd86 100644 --- a/ai2modulus/metrics/climate/reduction.py +++ b/modulus/metrics/climate/reduction.py @@ -15,7 +15,7 @@ import torch from torch import Tensor from abc import ABC -from ai2modulus.metrics.general.reduction import WeightedMean, WeightedVariance +from modulus.metrics.general.reduction import WeightedMean, WeightedVariance def _compute_lat_weights(lat: Tensor) -> Tensor: diff --git a/ai2modulus/metrics/general/__init__.py b/modulus/metrics/general/__init__.py similarity index 100% rename from ai2modulus/metrics/general/__init__.py rename to modulus/metrics/general/__init__.py diff --git a/ai2modulus/metrics/general/calibration.py b/modulus/metrics/general/calibration.py similarity index 98% rename from ai2modulus/metrics/general/calibration.py rename to modulus/metrics/general/calibration.py index 670837d3d4..ba96b50d86 100644 --- a/ai2modulus/metrics/general/calibration.py +++ b/modulus/metrics/general/calibration.py @@ -14,7 +14,7 @@ import torch import numpy as np -from ai2modulus.metrics.general.histogram import histogram, linspace +from modulus.metrics.general.histogram import histogram, linspace from typing import Union Tensor = torch.Tensor diff --git a/ai2modulus/metrics/general/crps.py b/modulus/metrics/general/crps.py similarity index 100% rename from ai2modulus/metrics/general/crps.py rename to modulus/metrics/general/crps.py diff --git a/ai2modulus/metrics/general/ensemble_metrics.py b/modulus/metrics/general/ensemble_metrics.py similarity index 99% rename from ai2modulus/metrics/general/ensemble_metrics.py rename to modulus/metrics/general/ensemble_metrics.py index 035de5780c..2258bf02fb 100644 --- a/ai2modulus/metrics/general/ensemble_metrics.py +++ b/modulus/metrics/general/ensemble_metrics.py @@ -16,7 +16,7 @@ import torch.distributed as dist from typing import Union, Tuple, List from abc import ABC, abstractmethod -from ai2modulus.distributed.manager import DistributedManager +from modulus.distributed.manager import DistributedManager from warnings import warn Tensor = torch.Tensor diff --git a/ai2modulus/metrics/general/entropy.py b/modulus/metrics/general/entropy.py similarity index 100% rename from ai2modulus/metrics/general/entropy.py rename to modulus/metrics/general/entropy.py diff --git a/ai2modulus/metrics/general/histogram.py b/modulus/metrics/general/histogram.py similarity index 99% rename from ai2modulus/metrics/general/histogram.py rename to modulus/metrics/general/histogram.py index 3eaf64979c..89564a7ac9 100644 --- a/ai2modulus/metrics/general/histogram.py +++ b/modulus/metrics/general/histogram.py @@ -17,7 +17,7 @@ import numpy import torch.distributed as dist from typing import Union, Tuple -from ai2modulus.distributed.manager import DistributedManager +from modulus.distributed.manager import DistributedManager from .ensemble_metrics import EnsembleMetrics Tensor = torch.Tensor diff --git a/ai2modulus/metrics/general/mse.py b/modulus/metrics/general/mse.py similarity index 100% rename from ai2modulus/metrics/general/mse.py rename to modulus/metrics/general/mse.py diff --git a/ai2modulus/metrics/general/reduction.py b/modulus/metrics/general/reduction.py similarity index 100% rename from ai2modulus/metrics/general/reduction.py rename to modulus/metrics/general/reduction.py diff --git a/ai2modulus/metrics/general/wasserstein.py b/modulus/metrics/general/wasserstein.py similarity index 100% rename from ai2modulus/metrics/general/wasserstein.py rename to modulus/metrics/general/wasserstein.py diff --git a/ai2modulus/models/__init__.py b/modulus/models/__init__.py similarity index 100% rename from ai2modulus/models/__init__.py rename to modulus/models/__init__.py diff --git a/ai2modulus/models/afno/__init__.py b/modulus/models/afno/__init__.py similarity index 100% rename from ai2modulus/models/afno/__init__.py rename to modulus/models/afno/__init__.py diff --git a/ai2modulus/models/afno/afno.py b/modulus/models/afno/afno.py similarity index 99% rename from ai2modulus/models/afno/afno.py rename to modulus/models/afno/afno.py index ab1810c2e1..619c543619 100644 --- a/ai2modulus/models/afno/afno.py +++ b/modulus/models/afno/afno.py @@ -15,8 +15,8 @@ import torch import torch.nn as nn import torch.nn.functional as F -import ai2modulus -import ai2modulus.models.layers.fft as fft +import modulus +import modulus.models.layers.fft as fft from functools import partial from typing import Tuple, Any diff --git a/ai2modulus/models/afno/distributed/__init__.py b/modulus/models/afno/distributed/__init__.py similarity index 100% rename from ai2modulus/models/afno/distributed/__init__.py rename to modulus/models/afno/distributed/__init__.py diff --git a/ai2modulus/models/afno/distributed/afno.py b/modulus/models/afno/distributed/afno.py similarity index 100% rename from ai2modulus/models/afno/distributed/afno.py rename to modulus/models/afno/distributed/afno.py diff --git a/ai2modulus/models/afno/distributed/layers.py b/modulus/models/afno/distributed/layers.py similarity index 100% rename from ai2modulus/models/afno/distributed/layers.py rename to modulus/models/afno/distributed/layers.py diff --git a/ai2modulus/models/afno/distributed/mappings.py b/modulus/models/afno/distributed/mappings.py similarity index 100% rename from ai2modulus/models/afno/distributed/mappings.py rename to modulus/models/afno/distributed/mappings.py diff --git a/ai2modulus/models/fno/__init__.py b/modulus/models/fno/__init__.py similarity index 100% rename from ai2modulus/models/fno/__init__.py rename to modulus/models/fno/__init__.py diff --git a/ai2modulus/models/fno/fno.py b/modulus/models/fno/fno.py similarity index 99% rename from ai2modulus/models/fno/fno.py rename to modulus/models/fno/fno.py index 13334785b0..1e993d62df 100644 --- a/ai2modulus/models/fno/fno.py +++ b/modulus/models/fno/fno.py @@ -15,7 +15,8 @@ import torch import torch.nn as nn import torch.nn.functional as F -import ai2modulus.models.layers as layers +import modulus.models.layers as layers +import modulus from typing import Dict, List, Union, Tuple from torch import Tensor @@ -484,7 +485,7 @@ class FNO(Module): Parameters ---------- - decoder_net : ai2modulus.Module + decoder_net : modulus.Module Pointwise decoder network, input feature size should match `latent_channels` in_channels : int Number of input channels @@ -508,14 +509,14 @@ class FNO(Module): Example ------- >>> # define the decoder net - >>> decoder = ai2modulus.models.mlp.FullyConnected( + >>> decoder = modulus.models.mlp.FullyConnected( ... in_features=32, ... out_features=3, ... num_layers=2, ... layer_size=16, ... ) >>> # define the 2d FNO model - >>> model = ai2modulus.models.fno.FNO( + >>> model = modulus.models.fno.FNO( ... decoder_net=decoder, ... in_channels=4, ... dimension=2, diff --git a/ai2modulus/models/gnn_layers/__init__.py b/modulus/models/gnn_layers/__init__.py similarity index 100% rename from ai2modulus/models/gnn_layers/__init__.py rename to modulus/models/gnn_layers/__init__.py diff --git a/ai2modulus/models/gnn_layers/embedder.py b/modulus/models/gnn_layers/embedder.py similarity index 100% rename from ai2modulus/models/gnn_layers/embedder.py rename to modulus/models/gnn_layers/embedder.py diff --git a/ai2modulus/models/gnn_layers/mesh_edge_block.py b/modulus/models/gnn_layers/mesh_edge_block.py similarity index 100% rename from ai2modulus/models/gnn_layers/mesh_edge_block.py rename to modulus/models/gnn_layers/mesh_edge_block.py diff --git a/ai2modulus/models/gnn_layers/mesh_graph_decoder.py b/modulus/models/gnn_layers/mesh_graph_decoder.py similarity index 100% rename from ai2modulus/models/gnn_layers/mesh_graph_decoder.py rename to modulus/models/gnn_layers/mesh_graph_decoder.py diff --git a/ai2modulus/models/gnn_layers/mesh_graph_encoder.py b/modulus/models/gnn_layers/mesh_graph_encoder.py similarity index 100% rename from ai2modulus/models/gnn_layers/mesh_graph_encoder.py rename to modulus/models/gnn_layers/mesh_graph_encoder.py diff --git a/ai2modulus/models/gnn_layers/mesh_graph_mlp.py b/modulus/models/gnn_layers/mesh_graph_mlp.py similarity index 99% rename from ai2modulus/models/gnn_layers/mesh_graph_mlp.py rename to modulus/models/gnn_layers/mesh_graph_mlp.py index eac38742f1..c0f94a69ff 100644 --- a/ai2modulus/models/gnn_layers/mesh_graph_mlp.py +++ b/modulus/models/gnn_layers/mesh_graph_mlp.py @@ -22,7 +22,7 @@ from torch.autograd.function import once_differentiable from .utils import concat_efeat, sum_efeat, CuGraphCSC -from ai2modulus.models.layers.fused_silu import silu_backward_for +from modulus.models.layers.fused_silu import silu_backward_for try: from apex.normalization import FusedLayerNorm diff --git a/ai2modulus/models/gnn_layers/mesh_node_block.py b/modulus/models/gnn_layers/mesh_node_block.py similarity index 100% rename from ai2modulus/models/gnn_layers/mesh_node_block.py rename to modulus/models/gnn_layers/mesh_node_block.py diff --git a/ai2modulus/models/gnn_layers/utils.py b/modulus/models/gnn_layers/utils.py similarity index 100% rename from ai2modulus/models/gnn_layers/utils.py rename to modulus/models/gnn_layers/utils.py diff --git a/ai2modulus/models/graphcast/__init__.py b/modulus/models/graphcast/__init__.py similarity index 100% rename from ai2modulus/models/graphcast/__init__.py rename to modulus/models/graphcast/__init__.py diff --git a/ai2modulus/models/graphcast/graph_cast_net.py b/modulus/models/graphcast/graph_cast_net.py similarity index 97% rename from ai2modulus/models/graphcast/graph_cast_net.py rename to modulus/models/graphcast/graph_cast_net.py index e3d771f64d..46f87759f0 100644 --- a/ai2modulus/models/graphcast/graph_cast_net.py +++ b/modulus/models/graphcast/graph_cast_net.py @@ -19,18 +19,18 @@ from typing import Any from dataclasses import dataclass -from ai2modulus.models.gnn_layers.utils import set_checkpoint_fn, CuGraphCSC -from ai2modulus.models.gnn_layers.embedder import ( +from modulus.models.gnn_layers.utils import set_checkpoint_fn, CuGraphCSC +from modulus.models.gnn_layers.embedder import ( GraphCastEncoderEmbedder, GraphCastDecoderEmbedder, ) -from ai2modulus.models.gnn_layers.mesh_graph_encoder import MeshGraphEncoder -from ai2modulus.models.gnn_layers.mesh_graph_decoder import MeshGraphDecoder -from ai2modulus.models.gnn_layers.mesh_graph_mlp import MeshGraphMLP -from ai2modulus.models.module import Module -from ai2modulus.models.meta import ModelMetaData -from ai2modulus.utils.graphcast.graph import Graph -from ai2modulus.utils.graphcast.data_utils import StaticData +from modulus.models.gnn_layers.mesh_graph_encoder import MeshGraphEncoder +from modulus.models.gnn_layers.mesh_graph_decoder import MeshGraphDecoder +from modulus.models.gnn_layers.mesh_graph_mlp import MeshGraphMLP +from modulus.models.module import Module +from modulus.models.meta import ModelMetaData +from modulus.utils.graphcast.graph import Graph +from modulus.utils.graphcast.data_utils import StaticData from .graph_cast_processor import GraphCastProcessor diff --git a/ai2modulus/models/graphcast/graph_cast_processor.py b/modulus/models/graphcast/graph_cast_processor.py similarity index 96% rename from ai2modulus/models/graphcast/graph_cast_processor.py rename to modulus/models/graphcast/graph_cast_processor.py index 8ac4d785ca..763bb11c1a 100644 --- a/ai2modulus/models/graphcast/graph_cast_processor.py +++ b/modulus/models/graphcast/graph_cast_processor.py @@ -19,9 +19,9 @@ from torch import Tensor from dgl import DGLGraph -from ai2modulus.models.gnn_layers.utils import set_checkpoint_fn, CuGraphCSC -from ai2modulus.models.gnn_layers.mesh_node_block import MeshNodeBlock -from ai2modulus.models.gnn_layers.mesh_edge_block import MeshEdgeBlock +from modulus.models.gnn_layers.utils import set_checkpoint_fn, CuGraphCSC +from modulus.models.gnn_layers.mesh_node_block import MeshNodeBlock +from modulus.models.gnn_layers.mesh_edge_block import MeshEdgeBlock class GraphCastProcessor(nn.Module): diff --git a/ai2modulus/models/layers/__init__.py b/modulus/models/layers/__init__.py similarity index 100% rename from ai2modulus/models/layers/__init__.py rename to modulus/models/layers/__init__.py diff --git a/ai2modulus/models/layers/activations.py b/modulus/models/layers/activations.py similarity index 99% rename from ai2modulus/models/layers/activations.py rename to modulus/models/layers/activations.py index 1ea2a87fa0..63ce2054c2 100644 --- a/ai2modulus/models/layers/activations.py +++ b/modulus/models/layers/activations.py @@ -14,7 +14,7 @@ import torch import torch.nn as nn -import ai2modulus +import modulus Tensor = torch.Tensor diff --git a/ai2modulus/models/layers/fft.py b/modulus/models/layers/fft.py similarity index 100% rename from ai2modulus/models/layers/fft.py rename to modulus/models/layers/fft.py diff --git a/ai2modulus/models/layers/fully_connected_layers.py b/modulus/models/layers/fully_connected_layers.py similarity index 100% rename from ai2modulus/models/layers/fully_connected_layers.py rename to modulus/models/layers/fully_connected_layers.py diff --git a/ai2modulus/models/layers/fused_silu.py b/modulus/models/layers/fused_silu.py similarity index 100% rename from ai2modulus/models/layers/fused_silu.py rename to modulus/models/layers/fused_silu.py diff --git a/ai2modulus/models/layers/spectral_layers.py b/modulus/models/layers/spectral_layers.py similarity index 100% rename from ai2modulus/models/layers/spectral_layers.py rename to modulus/models/layers/spectral_layers.py diff --git a/ai2modulus/models/layers/weight_norm.py b/modulus/models/layers/weight_norm.py similarity index 99% rename from ai2modulus/models/layers/weight_norm.py rename to modulus/models/layers/weight_norm.py index ef59a5d901..f341a4b686 100644 --- a/ai2modulus/models/layers/weight_norm.py +++ b/modulus/models/layers/weight_norm.py @@ -15,7 +15,7 @@ import torch import torch.nn as nn import torch.nn.functional as F -import ai2modulus +import modulus Tensor = torch.Tensor diff --git a/ai2modulus/models/meshgraphnet/__init__.py b/modulus/models/meshgraphnet/__init__.py similarity index 100% rename from ai2modulus/models/meshgraphnet/__init__.py rename to modulus/models/meshgraphnet/__init__.py diff --git a/ai2modulus/models/meshgraphnet/meshgraphnet.py b/modulus/models/meshgraphnet/meshgraphnet.py similarity index 95% rename from ai2modulus/models/meshgraphnet/meshgraphnet.py rename to modulus/models/meshgraphnet/meshgraphnet.py index 74a9b75428..2b52f09de0 100644 --- a/ai2modulus/models/meshgraphnet/meshgraphnet.py +++ b/modulus/models/meshgraphnet/meshgraphnet.py @@ -15,7 +15,7 @@ import torch from torch import Tensor import torch.nn as nn -import ai2modulus +import modulus try: import dgl @@ -29,13 +29,13 @@ from typing import Union, List from dataclasses import dataclass -from ai2modulus.models.meta import ModelMetaData -from ai2modulus.models.module import Module +from modulus.models.meta import ModelMetaData +from modulus.models.module import Module -from ai2modulus.models.gnn_layers.utils import CuGraphCSC -from ai2modulus.models.gnn_layers.mesh_graph_mlp import MeshGraphMLP -from ai2modulus.models.gnn_layers.mesh_edge_block import MeshEdgeBlock -from ai2modulus.models.gnn_layers.mesh_node_block import MeshNodeBlock +from modulus.models.gnn_layers.utils import CuGraphCSC +from modulus.models.gnn_layers.mesh_graph_mlp import MeshGraphMLP +from modulus.models.gnn_layers.mesh_edge_block import MeshEdgeBlock +from modulus.models.gnn_layers.mesh_node_block import MeshNodeBlock @dataclass diff --git a/ai2modulus/models/meta.py b/modulus/models/meta.py similarity index 100% rename from ai2modulus/models/meta.py rename to modulus/models/meta.py diff --git a/ai2modulus/models/mlp/__init__.py b/modulus/models/mlp/__init__.py similarity index 100% rename from ai2modulus/models/mlp/__init__.py rename to modulus/models/mlp/__init__.py diff --git a/ai2modulus/models/mlp/fully_connected.py b/modulus/models/mlp/fully_connected.py similarity index 98% rename from ai2modulus/models/mlp/fully_connected.py rename to modulus/models/mlp/fully_connected.py index 83603e599e..2bfd91043e 100644 --- a/ai2modulus/models/mlp/fully_connected.py +++ b/modulus/models/mlp/fully_connected.py @@ -14,12 +14,12 @@ import torch import torch.nn as nn -import ai2modulus +import modulus from torch import Tensor from dataclasses import dataclass from typing import Optional, Union, List -from ai2modulus.models.layers import FCLayer +from modulus.models.layers import FCLayer from ..meta import ModelMetaData from ..module import Module diff --git a/ai2modulus/models/module.py b/modulus/models/module.py similarity index 98% rename from ai2modulus/models/module.py rename to modulus/models/module.py index f29de35698..ffe10a0b3e 100644 --- a/ai2modulus/models/module.py +++ b/modulus/models/module.py @@ -17,7 +17,7 @@ from typing import Union from pathlib import Path -from ai2modulus.models.meta import ModelMetaData +from modulus.models.meta import ModelMetaData class Module(torch.nn.Module): diff --git a/ai2modulus/models/pix2pix/__init__.py b/modulus/models/pix2pix/__init__.py similarity index 100% rename from ai2modulus/models/pix2pix/__init__.py rename to modulus/models/pix2pix/__init__.py diff --git a/ai2modulus/models/pix2pix/pix2pix.py b/modulus/models/pix2pix/pix2pix.py similarity index 99% rename from ai2modulus/models/pix2pix/pix2pix.py rename to modulus/models/pix2pix/pix2pix.py index 12b74f88b9..d4bc837408 100644 --- a/ai2modulus/models/pix2pix/pix2pix.py +++ b/modulus/models/pix2pix/pix2pix.py @@ -58,7 +58,7 @@ from typing import Union, List from dataclasses import dataclass -import ai2modulus +import modulus from ..meta import ModelMetaData from ..module import Module diff --git a/ai2modulus/models/rnn/__init__.py b/modulus/models/rnn/__init__.py similarity index 100% rename from ai2modulus/models/rnn/__init__.py rename to modulus/models/rnn/__init__.py diff --git a/ai2modulus/models/rnn/layers.py b/modulus/models/rnn/layers.py similarity index 100% rename from ai2modulus/models/rnn/layers.py rename to modulus/models/rnn/layers.py diff --git a/ai2modulus/models/rnn/rnn_one2many.py b/modulus/models/rnn/rnn_one2many.py similarity index 98% rename from ai2modulus/models/rnn/rnn_one2many.py rename to modulus/models/rnn/rnn_one2many.py index 9ca00032a8..8288afb017 100644 --- a/ai2modulus/models/rnn/rnn_one2many.py +++ b/modulus/models/rnn/rnn_one2many.py @@ -14,14 +14,14 @@ import torch import torch.nn as nn -import ai2modulus +import modulus from torch import Tensor from dataclasses import dataclass from typing import Union, List -from ai2modulus.models.meta import ModelMetaData -from ai2modulus.models.module import Module -from ai2modulus.models.rnn.layers import ( +from modulus.models.meta import ModelMetaData +from modulus.models.module import Module +from modulus.models.rnn.layers import ( _ConvLayer, _TransposeConvLayer, _ConvGRULayer, diff --git a/ai2modulus/models/rnn/rnn_seq2seq.py b/modulus/models/rnn/rnn_seq2seq.py similarity index 98% rename from ai2modulus/models/rnn/rnn_seq2seq.py rename to modulus/models/rnn/rnn_seq2seq.py index 0a176fd459..97dad97129 100644 --- a/ai2modulus/models/rnn/rnn_seq2seq.py +++ b/modulus/models/rnn/rnn_seq2seq.py @@ -14,14 +14,14 @@ import torch import torch.nn as nn -import ai2modulus +import modulus from torch import Tensor from dataclasses import dataclass from typing import Union, List -from ai2modulus.models.meta import ModelMetaData -from ai2modulus.models.module import Module -from ai2modulus.models.rnn.layers import ( +from modulus.models.meta import ModelMetaData +from modulus.models.module import Module +from modulus.models.rnn.layers import ( _ConvLayer, _TransposeConvLayer, _ConvGRULayer, diff --git a/ai2modulus/models/sfno/activations.py b/modulus/models/sfno/activations.py similarity index 100% rename from ai2modulus/models/sfno/activations.py rename to modulus/models/sfno/activations.py diff --git a/ai2modulus/models/sfno/contractions.py b/modulus/models/sfno/contractions.py similarity index 100% rename from ai2modulus/models/sfno/contractions.py rename to modulus/models/sfno/contractions.py diff --git a/ai2modulus/models/sfno/factorizations.py b/modulus/models/sfno/factorizations.py similarity index 99% rename from ai2modulus/models/sfno/factorizations.py rename to modulus/models/sfno/factorizations.py index 37e757a5a4..48fd2ec3e3 100644 --- a/ai2modulus/models/sfno/factorizations.py +++ b/modulus/models/sfno/factorizations.py @@ -20,7 +20,7 @@ # from tensorly.plugins import use_opt_einsum # use_opt_einsum('optimal') -from ai2modulus.models.sfno.contractions import ( +from modulus.models.sfno.contractions import ( _contract_diagonal, _contract_dhconv, _contract_sep_diagonal, diff --git a/ai2modulus/models/sfno/initialization.py b/modulus/models/sfno/initialization.py similarity index 100% rename from ai2modulus/models/sfno/initialization.py rename to modulus/models/sfno/initialization.py diff --git a/ai2modulus/models/sfno/layers.py b/modulus/models/sfno/layers.py similarity index 99% rename from ai2modulus/models/sfno/layers.py rename to modulus/models/sfno/layers.py index 3a5338f412..85642c0e92 100644 --- a/ai2modulus/models/sfno/layers.py +++ b/modulus/models/sfno/layers.py @@ -21,9 +21,9 @@ import math from torch_harmonics import * -from ai2modulus.models.sfno.contractions import * -from ai2modulus.models.sfno.activations import * -from ai2modulus.models.sfno.initialization import trunc_normal_ +from modulus.models.sfno.contractions import * +from modulus.models.sfno.activations import * +from modulus.models.sfno.initialization import trunc_normal_ @torch.jit.script diff --git a/ai2modulus/models/sfno/preprocessor.py b/modulus/models/sfno/preprocessor.py similarity index 100% rename from ai2modulus/models/sfno/preprocessor.py rename to modulus/models/sfno/preprocessor.py diff --git a/ai2modulus/models/sfno/s2convolutions.py b/modulus/models/sfno/s2convolutions.py similarity index 98% rename from ai2modulus/models/sfno/s2convolutions.py rename to modulus/models/sfno/s2convolutions.py index 85406e6c0c..f53e725864 100644 --- a/ai2modulus/models/sfno/s2convolutions.py +++ b/modulus/models/sfno/s2convolutions.py @@ -27,13 +27,13 @@ from tltorch.factorized_tensors.core import FactorizedTensor # import convenience functions for factorized tensors -from ai2modulus.models.sfno.activations import ComplexReLU -from ai2modulus.models.sfno.contractions import compl_muladd2d_fwd, compl_mul2d_fwd -from ai2modulus.models.sfno.contractions import _contract_localconv_fwd -from ai2modulus.models.sfno.factorizations import get_contract_fun +from modulus.models.sfno.activations import ComplexReLU +from modulus.models.sfno.contractions import compl_muladd2d_fwd, compl_mul2d_fwd +from modulus.models.sfno.contractions import _contract_localconv_fwd +from modulus.models.sfno.factorizations import get_contract_fun # for the experimental module -from ai2modulus.models.sfno.contractions import ( +from modulus.models.sfno.contractions import ( compl_exp_muladd2d_fwd, compl_exp_mul2d_fwd, real_mul2d_fwd, diff --git a/ai2modulus/models/sfno/sfnonet.py b/modulus/models/sfno/sfnonet.py similarity index 98% rename from ai2modulus/models/sfno/sfnonet.py rename to modulus/models/sfno/sfnonet.py index 657332cc8d..7b0b75b9a5 100644 --- a/ai2modulus/models/sfno/sfnonet.py +++ b/modulus/models/sfno/sfnonet.py @@ -20,26 +20,26 @@ from typing import Any, Tuple # helpers -from ai2modulus.models.sfno.layers import trunc_normal_, DropPath, MLP +from modulus.models.sfno.layers import trunc_normal_, DropPath, MLP # import global convolution and non-linear spectral layers -from ai2modulus.models.sfno.layers import SpectralAttention2d -from ai2modulus.models.sfno.s2convolutions import SpectralConvS2, SpectralAttentionS2 +from modulus.models.sfno.layers import SpectralAttention2d +from modulus.models.sfno.s2convolutions import SpectralConvS2, SpectralAttentionS2 # get spectral transforms from torch_harmonics import torch_harmonics as th import torch_harmonics.distributed as thd # wrap fft, to unify interface to spectral transforms -from ai2modulus.models.sfno.layers import RealFFT2 -from ai2modulus.utils.sfno.distributed.layers import ( +from modulus.models.sfno.layers import RealFFT2 +from modulus.utils.sfno.distributed.layers import ( DistributedRealFFT2, DistributedInverseRealFFT2, DistributedMLP, ) # more distributed stuff -from ai2modulus.utils.sfno.distributed import comm +from modulus.utils.sfno.distributed import comm # layer normalization try: @@ -47,10 +47,10 @@ apex_imported = True except ImportError: apex_imported = False -from ai2modulus.utils.sfno.distributed.layer_norm import DistributedInstanceNorm2d +from modulus.utils.sfno.distributed.layer_norm import DistributedInstanceNorm2d -from ai2modulus.models.module import Module -from ai2modulus.models.meta import ModelMetaData +from modulus.models.module import Module +from modulus.models.meta import ModelMetaData @dataclass diff --git a/ai2modulus/models/srrn/__init__.py b/modulus/models/srrn/__init__.py similarity index 100% rename from ai2modulus/models/srrn/__init__.py rename to modulus/models/srrn/__init__.py diff --git a/ai2modulus/models/srrn/super_res_net.py b/modulus/models/srrn/super_res_net.py similarity index 99% rename from ai2modulus/models/srrn/super_res_net.py rename to modulus/models/srrn/super_res_net.py index ce339219e9..cdc1a210ee 100644 --- a/ai2modulus/models/srrn/super_res_net.py +++ b/modulus/models/srrn/super_res_net.py @@ -36,7 +36,7 @@ from typing import Union, List, Any from dataclasses import dataclass -import ai2modulus +import modulus from ..meta import ModelMetaData from ..module import Module diff --git a/ai2modulus/utils/__init__.py b/modulus/utils/__init__.py similarity index 100% rename from ai2modulus/utils/__init__.py rename to modulus/utils/__init__.py diff --git a/ai2modulus/utils/capture.py b/modulus/utils/capture.py similarity index 95% rename from ai2modulus/utils/capture.py rename to modulus/utils/capture.py index c1f906d4ca..9ca2a9b963 100644 --- a/ai2modulus/utils/capture.py +++ b/modulus/utils/capture.py @@ -13,7 +13,7 @@ # limitations under the License. import functools -import ai2modulus +import modulus import torch import logging from logging import Logger @@ -39,7 +39,7 @@ class _StaticCapture(object): def __init__( self, - model: ai2modulus.Module, + model: modulus.Module, optim: Union[optim, None] = None, logger: Union[Logger, None] = None, use_graphs: bool = True, @@ -52,10 +52,10 @@ def __init__( self.logger = logging.getLogger("capture") # DDP fix - if not isinstance(model, ai2modulus.Module) and hasattr(model, "module"): + if not isinstance(model, modulus.Module) and hasattr(model, "module"): model = model.module - if not isinstance(model, ai2modulus.Module): + if not isinstance(model, modulus.Module): self.logger.error("Model not a Modulus Module!") raise ValueError("Model not a Modulus Module!") self.model = model @@ -266,7 +266,7 @@ class StaticCaptureTraining(_StaticCapture): Parameters ---------- - model : ai2modulus.Module + model : modulus.Module Modulus Model optim : torch.optim Optimizer @@ -284,12 +284,12 @@ class StaticCaptureTraining(_StaticCapture): Raises ------ ValueError - If the model provided is not a ai2modulus.Module. I.e. has no meta data. + If the model provided is not a modulus.Module. I.e. has no meta data. Example ------- >>> # Create model - >>> model = ai2modulus.models.mlp.FullyConnected(2, 64, 2) + >>> model = modulus.models.mlp.FullyConnected(2, 64, 2) >>> input = torch.rand(8, 2) >>> output = torch.rand(8, 2) >>> # Create optimizer @@ -320,7 +320,7 @@ class StaticCaptureTraining(_StaticCapture): def __init__( self, - model: ai2modulus.Module, + model: modulus.Module, optim: torch.optim, logger: Union[Logger, None] = None, use_graphs: bool = True, @@ -349,7 +349,7 @@ class StaticCaptureEvaluateNoGrad(_StaticCapture): Parameters ---------- - model : ai2modulus.Module + model : modulus.Module Modulus Model logger : Union[Logger, None], optional Modulus Launch Logger, by default None @@ -365,12 +365,12 @@ class StaticCaptureEvaluateNoGrad(_StaticCapture): Raises ------ ValueError - If the model provided is not a ai2modulus.Module. I.e. has no meta data. + If the model provided is not a modulus.Module. I.e. has no meta data. Example ------- >>> # Create model - >>> model = ai2modulus.models.mlp.FullyConnected(2, 64, 2) + >>> model = modulus.models.mlp.FullyConnected(2, 64, 2) >>> input = torch.rand(8, 2) >>> # Create evaluate function with optimization wrapper >>> @StaticCaptureEvaluateNoGrad(model=model) @@ -391,7 +391,7 @@ class StaticCaptureEvaluateNoGrad(_StaticCapture): def __init__( self, - model: ai2modulus.Module, + model: modulus.Module, logger: Union[Logger, None] = None, use_graphs: bool = True, use_amp: bool = True, diff --git a/ai2modulus/utils/graphcast/__init__.py b/modulus/utils/graphcast/__init__.py similarity index 100% rename from ai2modulus/utils/graphcast/__init__.py rename to modulus/utils/graphcast/__init__.py diff --git a/ai2modulus/utils/graphcast/data_utils.py b/modulus/utils/graphcast/data_utils.py similarity index 100% rename from ai2modulus/utils/graphcast/data_utils.py rename to modulus/utils/graphcast/data_utils.py diff --git a/ai2modulus/utils/graphcast/graph.py b/modulus/utils/graphcast/graph.py similarity index 100% rename from ai2modulus/utils/graphcast/graph.py rename to modulus/utils/graphcast/graph.py diff --git a/ai2modulus/utils/graphcast/graph_utils.py b/modulus/utils/graphcast/graph_utils.py similarity index 100% rename from ai2modulus/utils/graphcast/graph_utils.py rename to modulus/utils/graphcast/graph_utils.py diff --git a/ai2modulus/utils/graphcast/icospheres.py b/modulus/utils/graphcast/icospheres.py similarity index 100% rename from ai2modulus/utils/graphcast/icospheres.py rename to modulus/utils/graphcast/icospheres.py diff --git a/ai2modulus/utils/graphcast/loss.py b/modulus/utils/graphcast/loss.py similarity index 100% rename from ai2modulus/utils/graphcast/loss.py rename to modulus/utils/graphcast/loss.py diff --git a/ai2modulus/utils/sfno/YParams.py b/modulus/utils/sfno/YParams.py similarity index 100% rename from ai2modulus/utils/sfno/YParams.py rename to modulus/utils/sfno/YParams.py diff --git a/ai2modulus/utils/sfno/distributed/__init__.py b/modulus/utils/sfno/distributed/__init__.py similarity index 100% rename from ai2modulus/utils/sfno/distributed/__init__.py rename to modulus/utils/sfno/distributed/__init__.py diff --git a/ai2modulus/utils/sfno/distributed/comm.py b/modulus/utils/sfno/distributed/comm.py similarity index 99% rename from ai2modulus/utils/sfno/distributed/comm.py rename to modulus/utils/sfno/distributed/comm.py index 10af570efe..e4da85db38 100644 --- a/ai2modulus/utils/sfno/distributed/comm.py +++ b/modulus/utils/sfno/distributed/comm.py @@ -14,7 +14,7 @@ import os import logging -from ai2modulus.utils.sfno.logging_utils import disable_logging +from modulus.utils.sfno.logging_utils import disable_logging import math import torch import torch.distributed as dist diff --git a/ai2modulus/utils/sfno/distributed/helpers.py b/modulus/utils/sfno/distributed/helpers.py similarity index 99% rename from ai2modulus/utils/sfno/distributed/helpers.py rename to modulus/utils/sfno/distributed/helpers.py index 27b545adae..d31cfc0f1d 100644 --- a/ai2modulus/utils/sfno/distributed/helpers.py +++ b/modulus/utils/sfno/distributed/helpers.py @@ -18,7 +18,7 @@ import torch.nn.functional as F import torch.distributed as dist -from ai2modulus.utils.sfno.distributed import comm +from modulus.utils.sfno.distributed import comm from torch._utils import _flatten_dense_tensors diff --git a/ai2modulus/utils/sfno/distributed/layer_norm.py b/modulus/utils/sfno/distributed/layer_norm.py similarity index 97% rename from ai2modulus/utils/sfno/distributed/layer_norm.py rename to modulus/utils/sfno/distributed/layer_norm.py index 052b0c9707..772ba7b50f 100644 --- a/ai2modulus/utils/sfno/distributed/layer_norm.py +++ b/modulus/utils/sfno/distributed/layer_norm.py @@ -20,8 +20,8 @@ from typing import Tuple # for spatial model-parallelism -from ai2modulus.utils.sfno.distributed import comm -from ai2modulus.utils.sfno.distributed.mappings import ( +from modulus.utils.sfno.distributed import comm +from modulus.utils.sfno.distributed.mappings import ( gather_from_parallel_region, copy_to_spatial_parallel_region, ) diff --git a/ai2modulus/utils/sfno/distributed/layers.py b/modulus/utils/sfno/distributed/layers.py similarity index 96% rename from ai2modulus/utils/sfno/distributed/layers.py rename to modulus/utils/sfno/distributed/layers.py index 9e5e432813..feafb9bed9 100644 --- a/ai2modulus/utils/sfno/distributed/layers.py +++ b/modulus/utils/sfno/distributed/layers.py @@ -15,21 +15,21 @@ import torch import torch.nn as nn import torch.nn.functional as F -from ai2modulus.utils.sfno.distributed import comm +from modulus.utils.sfno.distributed import comm # matmul parallel -from ai2modulus.utils.sfno.distributed.mappings import copy_to_matmul_parallel_region -from ai2modulus.utils.sfno.distributed.mappings import reduce_from_matmul_parallel_region -from ai2modulus.utils.sfno.distributed.mappings import scatter_to_matmul_parallel_region -from ai2modulus.utils.sfno.distributed.mappings import gather_from_matmul_parallel_region +from modulus.utils.sfno.distributed.mappings import copy_to_matmul_parallel_region +from modulus.utils.sfno.distributed.mappings import reduce_from_matmul_parallel_region +from modulus.utils.sfno.distributed.mappings import scatter_to_matmul_parallel_region +from modulus.utils.sfno.distributed.mappings import gather_from_matmul_parallel_region # spatial parallel -from ai2modulus.utils.sfno.distributed.mappings import gather_from_spatial_parallel_region -from ai2modulus.utils.sfno.distributed.mappings import scatter_to_spatial_parallel_region +from modulus.utils.sfno.distributed.mappings import gather_from_spatial_parallel_region +from modulus.utils.sfno.distributed.mappings import scatter_to_spatial_parallel_region -from ai2modulus.utils.sfno.distributed.helpers import _transpose +from modulus.utils.sfno.distributed.helpers import _transpose -from ai2modulus.models.sfno.initialization import trunc_normal_ +from modulus.models.sfno.initialization import trunc_normal_ class distributed_transpose_w(torch.autograd.Function): diff --git a/ai2modulus/utils/sfno/distributed/mappings.py b/modulus/utils/sfno/distributed/mappings.py similarity index 97% rename from ai2modulus/utils/sfno/distributed/mappings.py rename to modulus/utils/sfno/distributed/mappings.py index 725cb98791..5f41d652f1 100644 --- a/ai2modulus/utils/sfno/distributed/mappings.py +++ b/modulus/utils/sfno/distributed/mappings.py @@ -19,15 +19,15 @@ import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel -from ai2modulus.utils.sfno.distributed import comm +from modulus.utils.sfno.distributed import comm # torch utils from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors # helper functions -from ai2modulus.utils.sfno.distributed.helpers import _reduce -from ai2modulus.utils.sfno.distributed.helpers import _split -from ai2modulus.utils.sfno.distributed.helpers import _gather +from modulus.utils.sfno.distributed.helpers import _reduce +from modulus.utils.sfno.distributed.helpers import _split +from modulus.utils.sfno.distributed.helpers import _gather # generalized class _CopyToParallelRegion(torch.autograd.Function): diff --git a/ai2modulus/utils/sfno/img_utils.py b/modulus/utils/sfno/img_utils.py similarity index 100% rename from ai2modulus/utils/sfno/img_utils.py rename to modulus/utils/sfno/img_utils.py diff --git a/ai2modulus/utils/sfno/logging_utils.py b/modulus/utils/sfno/logging_utils.py similarity index 100% rename from ai2modulus/utils/sfno/logging_utils.py rename to modulus/utils/sfno/logging_utils.py diff --git a/ai2modulus/utils/sfno/loss.py b/modulus/utils/sfno/loss.py similarity index 99% rename from ai2modulus/utils/sfno/loss.py rename to modulus/utils/sfno/loss.py index 00b8afeb90..b32ffc12b0 100644 --- a/ai2modulus/utils/sfno/loss.py +++ b/modulus/utils/sfno/loss.py @@ -18,8 +18,8 @@ from torch import nn import torch.nn.functional as F -from ai2modulus.utils.sfno.distributed import comm -from ai2modulus.utils.sfno.distributed.mappings import gather_from_parallel_region +from modulus.utils.sfno.distributed import comm +from modulus.utils.sfno.distributed.mappings import gather_from_parallel_region import torch_harmonics as harmonics diff --git a/ai2modulus/utils/sfno/metric.py b/modulus/utils/sfno/metric.py similarity index 98% rename from ai2modulus/utils/sfno/metric.py rename to modulus/utils/sfno/metric.py index 9d5bd0b14c..67828e0174 100644 --- a/ai2modulus/utils/sfno/metric.py +++ b/modulus/utils/sfno/metric.py @@ -15,7 +15,7 @@ import torch # we need those here: -from ai2modulus.utils.sfno.metrics.weighted_acc_rmse import ( +from modulus.utils.sfno.metrics.weighted_acc_rmse import ( Quadrature, lat_torch, latitude_weighting_factor_torch, @@ -28,9 +28,9 @@ ) # distributed computing stuff -from ai2modulus.utils.sfno.distributed import comm +from modulus.utils.sfno.distributed import comm import torch.distributed as dist -from ai2modulus.utils.sfno.distributed.mappings import gather_from_parallel_region +from modulus.utils.sfno.distributed.mappings import gather_from_parallel_region class MetricsHandler: diff --git a/ai2modulus/utils/sfno/metrics/weighted_acc_rmse.py b/modulus/utils/sfno/metrics/weighted_acc_rmse.py similarity index 97% rename from ai2modulus/utils/sfno/metrics/weighted_acc_rmse.py rename to modulus/utils/sfno/metrics/weighted_acc_rmse.py index 2f83398108..141e9ee258 100644 --- a/ai2modulus/utils/sfno/metrics/weighted_acc_rmse.py +++ b/modulus/utils/sfno/metrics/weighted_acc_rmse.py @@ -13,7 +13,7 @@ # limitations under the License. import numpy as np -from ai2modulus.utils.sfno import logging_utils +from modulus.utils.sfno import logging_utils logging_utils.config_logger() import torch @@ -21,8 +21,8 @@ # distributed stuff -from ai2modulus.utils.sfno.distributed import comm -from ai2modulus.utils.sfno.distributed.mappings import reduce_from_matmul_parallel_region +from modulus.utils.sfno.distributed import comm +from modulus.utils.sfno.distributed.mappings import reduce_from_matmul_parallel_region def mean(x, axis=None): # pragma: no cover diff --git a/ai2modulus/utils/sfno/metrics/weighted_acc_rmse_inf.py b/modulus/utils/sfno/metrics/weighted_acc_rmse_inf.py similarity index 99% rename from ai2modulus/utils/sfno/metrics/weighted_acc_rmse_inf.py rename to modulus/utils/sfno/metrics/weighted_acc_rmse_inf.py index 1053e0a6df..b07223828b 100644 --- a/ai2modulus/utils/sfno/metrics/weighted_acc_rmse_inf.py +++ b/modulus/utils/sfno/metrics/weighted_acc_rmse_inf.py @@ -13,7 +13,7 @@ # limitations under the License. import numpy as np -from ai2modulus.utils.sfno import logging_utils +from modulus.utils.sfno import logging_utils logging_utils.config_logger() import torch diff --git a/ai2modulus/utils/sfno/warmup_scheduler.py b/modulus/utils/sfno/warmup_scheduler.py similarity index 100% rename from ai2modulus/utils/sfno/warmup_scheduler.py rename to modulus/utils/sfno/warmup_scheduler.py diff --git a/ai2modulus/utils/sfno/zenith_angle.py b/modulus/utils/sfno/zenith_angle.py similarity index 100% rename from ai2modulus/utils/sfno/zenith_angle.py rename to modulus/utils/sfno/zenith_angle.py diff --git a/pyproject.toml b/pyproject.toml index 6c111c10c7..26f7986e04 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ requires = ["setuptools", "setuptools-scm"] build-backend = "setuptools.build_meta" [project] -name = "nvidia-ai2-modulus" +name = "nvidia-modulus" authors = [ { name="NVIDIA Modulus Team"}, ] @@ -39,7 +39,7 @@ dynamic = ["version"] inference = ["onnxruntime==1.13.1"] [tool.setuptools.dynamic] -version = {attr = "ai2modulus.__version__"} +version = {attr = "modulus.__version__"} [tool.setuptools.packages.find] -include = ["ai2modulus", "ai2modulus.*"] +include = ["modulus", "modulus.*"] diff --git a/test/datapipes/common/cuda_graphs.py b/test/datapipes/common/cuda_graphs.py index 8f2c13edc5..9c66ceb836 100644 --- a/test/datapipes/common/cuda_graphs.py +++ b/test/datapipes/common/cuda_graphs.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import ai2modulus +import modulus import torch import logging @@ -56,7 +56,7 @@ def forward(self, inputs: Tuple[Tensor, ...]) -> Tuple[Tensor, ...]: def check_cuda_graphs( - datapipe: ai2modulus.Datapipe, + datapipe: modulus.Datapipe, input_fn: Union[Callable, None] = None, iterations: int = 5, warmup_length: int = 3, @@ -65,8 +65,8 @@ def check_cuda_graphs( Parameters ---------- - datapipe : ai2modulus.Datapipe - ai2modulus data pipe to test + datapipe : modulus.Datapipe + Modulus data pipe to test input_fn : Union[Callable, None], optional Input pre-processing function to produce a tuple of tensors for model inputs, by default None iterations : int, optional diff --git a/test/datapipes/common/device.py b/test/datapipes/common/device.py index cce4d87707..44aebb4147 100644 --- a/test/datapipes/common/device.py +++ b/test/datapipes/common/device.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import modulus + import torch import logging from typing import Tuple, Union diff --git a/test/datapipes/common/iterator_check.py b/test/datapipes/common/iterator_check.py index 34e5643d68..ce43a3daeb 100644 --- a/test/datapipes/common/iterator_check.py +++ b/test/datapipes/common/iterator_check.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import ai2modulus +import modulus import torch import logging @@ -20,7 +20,7 @@ logger = logging.getLogger("__name__") -def check_datapipe_iterable(datapipe: ai2modulus.Datapipe, nr_iterations: int = 3) -> bool: +def check_datapipe_iterable(datapipe: modulus.Datapipe, nr_iterations: int = 3) -> bool: """Checks if datapipe is iterable Parameters diff --git a/test/datapipes/common/sequence.py b/test/datapipes/common/sequence.py index 07955e6062..ae3d7f74bc 100644 --- a/test/datapipes/common/sequence.py +++ b/test/datapipes/common/sequence.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import modulus + import torch import logging from typing import Tuple, Union diff --git a/test/datapipes/common/shape_check.py b/test/datapipes/common/shape_check.py index 2d317bedb2..4ffead1188 100644 --- a/test/datapipes/common/shape_check.py +++ b/test/datapipes/common/shape_check.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import modulus import torch import logging diff --git a/test/datapipes/common/shuffle.py b/test/datapipes/common/shuffle.py index 60c010d084..8c2687d74d 100644 --- a/test/datapipes/common/shuffle.py +++ b/test/datapipes/common/shuffle.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import modulus + import torch import logging from typing import Tuple, Union diff --git a/test/datapipes/sfno/test_data_loader_dummy.py b/test/datapipes/sfno/test_data_loader_dummy.py index 13a907acf6..092c554a0d 100644 --- a/test/datapipes/sfno/test_data_loader_dummy.py +++ b/test/datapipes/sfno/test_data_loader_dummy.py @@ -16,7 +16,7 @@ import pytest from unittest.mock import MagicMock -from ai2modulus.datapipes.climate.sfno.dataloaders.data_loader_dummy import DummyLoader +from modulus.datapipes.climate.sfno.dataloaders.data_loader_dummy import DummyLoader def test_dummyloader(): diff --git a/test/datapipes/test_darcy.py b/test/datapipes/test_darcy.py index 950ee53d80..1d19b66608 100644 --- a/test/datapipes/test_darcy.py +++ b/test/datapipes/test_darcy.py @@ -16,7 +16,7 @@ import torch from typing import Tuple -from ai2modulus.datapipes.benchmarks.darcy import Darcy2D +from modulus.datapipes.benchmarks.darcy import Darcy2D from . import common Tensor = torch.Tensor diff --git a/test/datapipes/test_era5_hdf5.py b/test/datapipes/test_era5_hdf5.py index 7959b6e203..669675f4da 100644 --- a/test/datapipes/test_era5_hdf5.py +++ b/test/datapipes/test_era5_hdf5.py @@ -16,7 +16,7 @@ import torch from typing import Tuple -from ai2modulus.datapipes.climate import ERA5HDF5Datapipe +from modulus.datapipes.climate import ERA5HDF5Datapipe from . import common Tensor = torch.Tensor diff --git a/test/datapipes/test_kelvin_helmholtz.py b/test/datapipes/test_kelvin_helmholtz.py index df3d978465..b23e0505a3 100644 --- a/test/datapipes/test_kelvin_helmholtz.py +++ b/test/datapipes/test_kelvin_helmholtz.py @@ -16,7 +16,7 @@ import torch from typing import Tuple -from ai2modulus.datapipes.benchmarks.kelvin_helmholtz import KelvinHelmholtz2D +from modulus.datapipes.benchmarks.kelvin_helmholtz import KelvinHelmholtz2D from . import common Tensor = torch.Tensor diff --git a/test/deploy/test_onnx_fft.py b/test/deploy/test_onnx_fft.py index acf2a221c6..81affa5f99 100644 --- a/test/deploy/test_onnx_fft.py +++ b/test/deploy/test_onnx_fft.py @@ -19,7 +19,7 @@ import torch.nn as nn import torch.onnx import torch.onnx.utils -import ai2modulus.models.layers.fft as fft +import modulus.models.layers.fft as fft try: import onnxruntime as ort @@ -27,7 +27,7 @@ ort = None from typing import Tuple -from ai2modulus.deploy.onnx import export_to_onnx_stream, run_onnx_inference +from modulus.deploy.onnx import export_to_onnx_stream, run_onnx_inference Tensor = torch.Tensor logger = logging.getLogger("__name__") diff --git a/test/deploy/test_onnx_utils.py b/test/deploy/test_onnx_utils.py index 3e855cb82f..4cc0c921ca 100644 --- a/test/deploy/test_onnx_utils.py +++ b/test/deploy/test_onnx_utils.py @@ -23,8 +23,8 @@ ort = None from pathlib import Path -from ai2modulus.models.mlp import FullyConnected -from ai2modulus.deploy.onnx import export_to_onnx_stream, run_onnx_inference +from modulus.models.mlp import FullyConnected +from modulus.deploy.onnx import export_to_onnx_stream, run_onnx_inference Tensor = torch.Tensor logger = logging.getLogger("__name__") diff --git a/test/distributed/test_manager.py b/test/distributed/test_manager.py index f5ad5fa6ea..f4e25c23a8 100644 --- a/test/distributed/test_manager.py +++ b/test/distributed/test_manager.py @@ -13,7 +13,7 @@ # limitations under the License. import os -from ai2modulus.distributed import DistributedManager +from modulus.distributed import DistributedManager # TODO: Need to figure out how to test parallel set up def test_manager(): diff --git a/test/distributed/test_utils.py b/test/distributed/test_utils.py index 7836a7ace9..54ea4d7f46 100644 --- a/test/distributed/test_utils.py +++ b/test/distributed/test_utils.py @@ -13,8 +13,8 @@ # limitations under the License. import os -from ai2modulus.distributed import DistributedManager -from ai2modulus.distributed import gather_loss +from modulus.distributed import DistributedManager +from modulus.distributed import gather_loss # TODO: Need to figure out how to test parallel set up def test_gather_loss(): diff --git a/test/metrics/test_metrics_climate.py b/test/metrics/test_metrics_climate.py index a4a85920ef..842770f14a 100644 --- a/test/metrics/test_metrics_climate.py +++ b/test/metrics/test_metrics_climate.py @@ -15,13 +15,13 @@ import pytest import torch import numpy as np -from ai2modulus.metrics.general.mse import mse, rmse -from ai2modulus.metrics.climate.acc import acc -import ai2modulus.metrics.climate.reduction as clim_red -import ai2modulus.metrics.climate.efi as efi +from modulus.metrics.general.mse import mse, rmse +from modulus.metrics.climate.acc import acc +import modulus.metrics.climate.reduction as clim_red +import modulus.metrics.climate.efi as efi -import ai2modulus.metrics.general.reduction as gen_red -import ai2modulus.metrics.general.histogram as hist +import modulus.metrics.general.reduction as gen_red +import modulus.metrics.general.histogram as hist Tensor = torch.Tensor diff --git a/test/metrics/test_metrics_general.py b/test/metrics/test_metrics_general.py index 8dc3a0f8d4..3162b8433c 100644 --- a/test/metrics/test_metrics_general.py +++ b/test/metrics/test_metrics_general.py @@ -18,14 +18,14 @@ import numpy as np import os -import ai2modulus.metrics.general.histogram as hist -import ai2modulus.metrics.general.ensemble_metrics as em -import ai2modulus.metrics.general.crps as crps -import ai2modulus.metrics.general.wasserstein as w -import ai2modulus.metrics.general.calibration as cal -import ai2modulus.metrics.general.entropy as ent - -from ai2modulus.distributed.manager import DistributedManager +import modulus.metrics.general.histogram as hist +import modulus.metrics.general.ensemble_metrics as em +import modulus.metrics.general.crps as crps +import modulus.metrics.general.wasserstein as w +import modulus.metrics.general.calibration as cal +import modulus.metrics.general.entropy as ent + +from modulus.distributed.manager import DistributedManager Tensor = torch.Tensor diff --git a/test/models/common/checkpoints.py b/test/models/common/checkpoints.py index 6d69b5c9de..b1d67965c2 100644 --- a/test/models/common/checkpoints.py +++ b/test/models/common/checkpoints.py @@ -13,7 +13,7 @@ # limitations under the License. import logging -import ai2modulus +import modulus import torch from typing import Tuple @@ -26,8 +26,8 @@ @torch.no_grad() def validate_checkpoint( - model_1: ai2modulus.Module, - model_2: ai2modulus.Module, + model_1: modulus.Module, + model_2: modulus.Module, in_args: Tuple[Tensor] = (), rtol: float = 1e-5, atol: float = 1e-5, @@ -40,10 +40,10 @@ def validate_checkpoint( Parameters ---------- - model_1 : ai2modulus.Module - ai2modulus model to save checkpoint from - model_2 : ai2modulus.Module - ai2modulus model to load checkpoint to + model_1 : modulus.Module + Modulus model to save checkpoint from + model_2 : modulus.Module + Modulus model to load checkpoint to in_args : Tuple[Tensor], optional Input arguments, by default () rtol : float, optional diff --git a/test/models/common/fwdaccuracy.py b/test/models/common/fwdaccuracy.py index 9f57342d4d..d87d40ffe4 100644 --- a/test/models/common/fwdaccuracy.py +++ b/test/models/common/fwdaccuracy.py @@ -13,7 +13,7 @@ # limitations under the License. import logging -import ai2modulus +import modulus import torch import sys @@ -66,7 +66,7 @@ def save_output(output: Union[Tensor, Tuple[Tensor, ...]], file_name: Path): @torch.no_grad() def validate_forward_accuracy( - model: ai2modulus.Module, + model: modulus.Module, in_args: Tuple[Tensor] = (), rtol: float = 1e-3, atol: float = 1e-3, @@ -81,7 +81,7 @@ def validate_forward_accuracy( Parameters ---------- - model : ai2modulus.Module + model : modulus.Module Modulus module in_args : Tuple[Tensor], optional Input arguments, by default () diff --git a/test/models/common/inference.py b/test/models/common/inference.py index de4b107e01..f155bd14af 100644 --- a/test/models/common/inference.py +++ b/test/models/common/inference.py @@ -13,7 +13,7 @@ # limitations under the License. import logging -import ai2modulus +import modulus import torch import onnx import pytest @@ -25,7 +25,7 @@ from typing import Tuple from pathlib import Path -from ai2modulus.deploy.onnx import export_to_onnx_stream, run_onnx_inference +from modulus.deploy.onnx import export_to_onnx_stream, run_onnx_inference from .utils import compare_output Tensor = torch.Tensor @@ -50,7 +50,7 @@ def check_ort_version(): @torch.no_grad() def validate_onnx_export( - model: ai2modulus.Module, + model: modulus.Module, in_args: Tuple[Tensor] = (), ) -> bool: """Check network's ONNX export works @@ -59,7 +59,7 @@ def validate_onnx_export( Parameters ---------- - model_1 : ai2modulus.Module + model_1 : modulus.Module Modulus model to save checkpoint from in_args : Tuple[Tensor], optional Input arguments, by default () @@ -111,7 +111,7 @@ def validate_onnx_export( @torch.no_grad() def validate_onnx_runtime( - model: ai2modulus.Module, + model: modulus.Module, in_args: Tuple[Tensor, ...] = (), rtol: float = 1e-3, atol: float = 1e-3, @@ -125,7 +125,7 @@ def validate_onnx_runtime( Parameters ---------- - model_1 : ai2modulus.Module + model_1 : modulus.Module Modulus model to save checkpoint from in_args : Tuple[Tensor], optional Input arguments, by default () diff --git a/test/models/common/optimization.py b/test/models/common/optimization.py index 80441fe485..bef415e1d1 100644 --- a/test/models/common/optimization.py +++ b/test/models/common/optimization.py @@ -13,7 +13,7 @@ # limitations under the License. import logging -import ai2modulus +import modulus import torch from typing import Tuple @@ -27,7 +27,7 @@ @torch.no_grad() def validate_jit( - model: ai2modulus.Module, + model: modulus.Module, in_args: Tuple[Tensor] = (), rtol: float = 1e-5, atol: float = 1e-5, @@ -40,7 +40,7 @@ def validate_jit( Parameters ---------- - model : ai2modulus.Module + model : modulus.Module Modulus module in_args : Tuple[Tensor], optional Input arguments, by default () @@ -70,7 +70,7 @@ def validate_jit( def validate_cuda_graphs( - model: ai2modulus.Module, + model: modulus.Module, in_args: Tuple[Tensor] = (), rtol: float = 1e-5, atol: float = 1e-5, @@ -84,7 +84,7 @@ def validate_cuda_graphs( Parameters ---------- - model : ai2modulus.Module + model : modulus.Module Modulus module in_args : Tuple[Tensor], optional Input arguments, keywords not supported, by default () @@ -130,7 +130,7 @@ def validate_cuda_graphs( def validate_amp( - model: ai2modulus.Module, + model: modulus.Module, in_args: Tuple[Tensor] = (), iterations: int = 3, ) -> bool: @@ -140,7 +140,7 @@ def validate_amp( Parameters ---------- - model : ai2modulus.Module + model : modulus.Module Modulus module in_args : Tuple[Tensor], optional Input arguments, keywords not supported, by default () @@ -194,7 +194,7 @@ def validate_torch_fx() -> bool: def validate_combo_optims( - model: ai2modulus.Module, + model: modulus.Module, in_args: Tuple[Tensor] = (), iterations: int = 2, warmup_length: int = 11, @@ -207,7 +207,7 @@ def validate_combo_optims( Parameters ---------- - model : ai2modulus.Module + model : modulus.Module Modulus module in_args : Tuple[Tensor], optional Input arguments, keywords not supported, by default () diff --git a/test/models/graphcast/test_concat_trick.py b/test/models/graphcast/test_concat_trick.py index 6be56d6d42..4a2b119c6c 100644 --- a/test/models/graphcast/test_concat_trick.py +++ b/test/models/graphcast/test_concat_trick.py @@ -16,7 +16,7 @@ import numpy as np from utils import get_icosphere_path, fix_random_seeds -from ai2modulus.models.graphcast.graph_cast_net import GraphCastNet +from modulus.models.graphcast.graph_cast_net import GraphCastNet icosphere_path = get_icosphere_path() diff --git a/test/models/graphcast/test_cugraphops.py b/test/models/graphcast/test_cugraphops.py index 59b32ac145..6da949cdca 100644 --- a/test/models/graphcast/test_cugraphops.py +++ b/test/models/graphcast/test_cugraphops.py @@ -15,7 +15,7 @@ import torch import numpy as np from utils import get_icosphere_path, fix_random_seeds -from ai2modulus.models.graphcast.graph_cast_net import GraphCastNet +from modulus.models.graphcast.graph_cast_net import GraphCastNet def test_cugraphops(): diff --git a/test/models/graphcast/test_grad_checkpointing.py b/test/models/graphcast/test_grad_checkpointing.py index a5f0fffa7c..b43df6e133 100644 --- a/test/models/graphcast/test_grad_checkpointing.py +++ b/test/models/graphcast/test_grad_checkpointing.py @@ -16,7 +16,7 @@ import pytest from utils import fix_random_seeds, create_random_input, get_icosphere_path -from ai2modulus.models.graphcast.graph_cast_net import GraphCastNet +from modulus.models.graphcast.graph_cast_net import GraphCastNet @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) diff --git a/test/models/graphcast/test_graphcast.py b/test/models/graphcast/test_graphcast.py index db3145fe97..5a22c6ab7f 100644 --- a/test/models/graphcast/test_graphcast.py +++ b/test/models/graphcast/test_graphcast.py @@ -23,7 +23,7 @@ import common from utils import get_icosphere_path -from ai2modulus.models.graphcast.graph_cast_net import GraphCastNet +from modulus.models.graphcast.graph_cast_net import GraphCastNet icosphere_path = get_icosphere_path() diff --git a/test/models/sfno/test_activations.py b/test/models/sfno/test_activations.py index 65baabf64b..6ea069b261 100644 --- a/test/models/sfno/test_activations.py +++ b/test/models/sfno/test_activations.py @@ -16,7 +16,7 @@ import torch.nn.functional as F from torch import nn -from ai2modulus.models.sfno.activations import ComplexReLU, ComplexActivation +from modulus.models.sfno.activations import ComplexReLU, ComplexActivation def test_ComplexReLU_cartesian(): diff --git a/test/models/sfno/test_sfno.py b/test/models/sfno/test_sfno.py index 09b8dc40b4..4cf0aefaa7 100644 --- a/test/models/sfno/test_sfno.py +++ b/test/models/sfno/test_sfno.py @@ -22,7 +22,7 @@ import torch from utils import fix_random_seeds -from ai2modulus.models.sfno.sfnonet import SphericalFourierNeuralOperatorNet +from modulus.models.sfno.sfnonet import SphericalFourierNeuralOperatorNet @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) diff --git a/test/models/test_afno.py b/test/models/test_afno.py index 259da53d3b..c0265102c5 100644 --- a/test/models/test_afno.py +++ b/test/models/test_afno.py @@ -16,7 +16,7 @@ import pytest import random -from ai2modulus.models.afno import AFNO +from modulus.models.afno import AFNO from . import common diff --git a/test/models/test_fno.py b/test/models/test_fno.py index bd7d893a57..2f56ff8229 100644 --- a/test/models/test_fno.py +++ b/test/models/test_fno.py @@ -16,8 +16,8 @@ import pytest import random -from ai2modulus.models.mlp import FullyConnected -from ai2modulus.models.fno import FNO +from modulus.models.mlp import FullyConnected +from modulus.models.fno import FNO from . import common diff --git a/test/models/test_fully_connected.py b/test/models/test_fully_connected.py index f377aff65a..83ff62a637 100644 --- a/test/models/test_fully_connected.py +++ b/test/models/test_fully_connected.py @@ -16,7 +16,7 @@ import pytest import random -from ai2modulus.models.mlp import FullyConnected +from modulus.models.mlp import FullyConnected from . import common diff --git a/test/models/test_layers_activations.py b/test/models/test_layers_activations.py index 27097d9990..84478ea765 100644 --- a/test/models/test_layers_activations.py +++ b/test/models/test_layers_activations.py @@ -16,8 +16,8 @@ import pytest import random -from ai2modulus.models.layers.activations import Identity, Stan, SquarePlus -from ai2modulus.models.layers.fused_silu import ( +from modulus.models.layers.activations import Identity, Stan, SquarePlus +from modulus.models.layers.fused_silu import ( FusedSiLU, FusedSiLU_deriv_1, FusedSiLU_deriv_2, diff --git a/test/models/test_layers_weightnorm.py b/test/models/test_layers_weightnorm.py index fc7bb29eb8..4492c96020 100644 --- a/test/models/test_layers_weightnorm.py +++ b/test/models/test_layers_weightnorm.py @@ -16,7 +16,7 @@ import pytest import random -from ai2modulus.models.layers import WeightNormLinear +from modulus.models.layers import WeightNormLinear @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) diff --git a/test/models/test_meshgraphnet.py b/test/models/test_meshgraphnet.py index 8f8d48db79..7a4f1d3b19 100644 --- a/test/models/test_meshgraphnet.py +++ b/test/models/test_meshgraphnet.py @@ -18,7 +18,7 @@ import dgl import numpy as np -from ai2modulus.models.meshgraphnet import MeshGraphNet +from modulus.models.meshgraphnet import MeshGraphNet from . import common diff --git a/test/models/test_pix2pix.py b/test/models/test_pix2pix.py index a395972968..b7b5d815c8 100644 --- a/test/models/test_pix2pix.py +++ b/test/models/test_pix2pix.py @@ -16,7 +16,7 @@ import pytest import random -from ai2modulus.models.pix2pix import Pix2Pix +from modulus.models.pix2pix import Pix2Pix from . import common diff --git a/test/models/test_rnn.py b/test/models/test_rnn.py index 3923347306..20c275a97d 100644 --- a/test/models/test_rnn.py +++ b/test/models/test_rnn.py @@ -16,8 +16,8 @@ import pytest import random -from ai2modulus.models.rnn.rnn_one2many import One2ManyRNN -from ai2modulus.models.rnn.rnn_seq2seq import Seq2SeqRNN +from modulus.models.rnn.rnn_one2many import One2ManyRNN +from modulus.models.rnn.rnn_seq2seq import Seq2SeqRNN from . import common diff --git a/test/models/test_rnn_layers.py b/test/models/test_rnn_layers.py index fbb043392d..e62f04bc98 100644 --- a/test/models/test_rnn_layers.py +++ b/test/models/test_rnn_layers.py @@ -15,7 +15,7 @@ import torch import pytest -from ai2modulus.models.rnn.layers import ( +from modulus.models.rnn.layers import ( _ConvLayer, _TransposeConvLayer, _ConvResidualBlock, diff --git a/test/models/test_super_res_net.py b/test/models/test_super_res_net.py index 634d20f4ec..1b600db2bc 100644 --- a/test/models/test_super_res_net.py +++ b/test/models/test_super_res_net.py @@ -16,7 +16,7 @@ import pytest import random -from ai2modulus.models.srrn import SRResNet +from modulus.models.srrn import SRResNet from . import common diff --git a/test/utils/graphcast/test_coordinate_transform.py b/test/utils/graphcast/test_coordinate_transform.py index 9b01fc113f..eabee566c6 100644 --- a/test/utils/graphcast/test_coordinate_transform.py +++ b/test/utils/graphcast/test_coordinate_transform.py @@ -16,7 +16,7 @@ import torch import pytest -from ai2modulus.utils.graphcast.graph_utils import latlon2xyz, xyz2latlon +from modulus.utils.graphcast.graph_utils import latlon2xyz, xyz2latlon @pytest.mark.parametrize("latlon", [[-27.0, 48.0], [0, 0], [62.0, -45.0]]) diff --git a/test/utils/graphcast/test_loss.py b/test/utils/graphcast/test_loss.py index b14a654344..54b4b505f1 100644 --- a/test/utils/graphcast/test_loss.py +++ b/test/utils/graphcast/test_loss.py @@ -14,7 +14,7 @@ import torch -from ai2modulus.utils.graphcast.loss import ( +from modulus.utils.graphcast.loss import ( CustomCellAreaWeightedLossFunction, CellAreaWeightedLossFunction, ) diff --git a/test/utils/sfno/test_img_utils.py b/test/utils/sfno/test_img_utils.py index adacaa4805..02847f617a 100644 --- a/test/utils/sfno/test_img_utils.py +++ b/test/utils/sfno/test_img_utils.py @@ -15,7 +15,7 @@ import os import torch import numpy as np -from ai2modulus.utils.sfno.img_utils import PeriodicPad2d, reshape_fields +from modulus.utils.sfno.img_utils import PeriodicPad2d, reshape_fields def test_PeriodicPad2d(): diff --git a/test/utils/sfno/test_logging.py b/test/utils/sfno/test_logging.py index ea696b5467..5a219b9e4b 100644 --- a/test/utils/sfno/test_logging.py +++ b/test/utils/sfno/test_logging.py @@ -15,7 +15,7 @@ import logging from io import StringIO from contextlib import redirect_stdout -from ai2modulus.utils.sfno.logging_utils import config_logger, disable_logging +from modulus.utils.sfno.logging_utils import config_logger, disable_logging def test_disable_logging(): diff --git a/test/utils/sfno/test_warmup_scheduler.py b/test/utils/sfno/test_warmup_scheduler.py index 3b1844e96b..d55317dcb2 100644 --- a/test/utils/sfno/test_warmup_scheduler.py +++ b/test/utils/sfno/test_warmup_scheduler.py @@ -16,7 +16,7 @@ from torch import nn from torch.optim import lr_scheduler as lrs -from ai2modulus.utils.sfno.warmup_scheduler import WarmupScheduler +from modulus.utils.sfno.warmup_scheduler import WarmupScheduler def test_warmup_scheduler(): diff --git a/test/utils/sfno/test_yparams.py b/test/utils/sfno/test_yparams.py index 8184504028..c7eaf00927 100644 --- a/test/utils/sfno/test_yparams.py +++ b/test/utils/sfno/test_yparams.py @@ -13,7 +13,7 @@ # limitations under the License. import json -from ai2modulus.utils.sfno.YParams import ParamsBase, YParams +from modulus.utils.sfno.YParams import ParamsBase, YParams def test_ParamsBase(): diff --git a/test/utils/test_capture.py b/test/utils/test_capture.py index 293d2e97e9..13dbef783c 100644 --- a/test/utils/test_capture.py +++ b/test/utils/test_capture.py @@ -18,9 +18,9 @@ import torch import torch.nn as nn -from ai2modulus.models.mlp import FullyConnected -from ai2modulus.utils import StaticCaptureTraining, StaticCaptureEvaluateNoGrad -from ai2modulus.utils.capture import _StaticCapture +from modulus.models.mlp import FullyConnected +from modulus.utils import StaticCaptureTraining, StaticCaptureEvaluateNoGrad +from modulus.utils.capture import _StaticCapture try: from apex import optimizers