Move ntk computation to analysis

zincware · May 31, 2024 · 6e0fad9 · 6e0fad9
1 parent 3b5f5a1
commit 6e0fad9
Show file tree

Hide file tree

Showing 24 changed files with 30 additions and 60 deletions.
diff --git a/...nit_tests/ntk_computation/test_jax_ntk.py → CI/unit_tests/analysis/test_jax_ntk.py b/...nit_tests/ntk_computation/test_jax_ntk.py → CI/unit_tests/analysis/test_jax_ntk.py
@@ -31,8 +31,8 @@
 from flax import linen as nn
 from jax import random
 
+from znnl.analysis import JAXNTKComputation
 from znnl.models import FlaxModel
-from znnl.ntk_computation import JAXNTKComputation
 
 
 class FlaxTestModule(nn.Module):

diff --git a/...ntk_computation/test_jax_ntk_classwise.py → ..._tests/analysis/test_jax_ntk_classwise.py b/...ntk_computation/test_jax_ntk_classwise.py → ..._tests/analysis/test_jax_ntk_classwise.py
@@ -30,8 +30,8 @@
 from flax import linen as nn
 from jax import random
 
+from znnl.analysis import JAXNTKClassWise
 from znnl.models import FlaxModel
-from znnl.ntk_computation import JAXNTKClassWise
 
 
 class FlaxTestModule(nn.Module):

diff --git a/..._computation/test_jax_ntk_combinations.py → ...sts/analysis/test_jax_ntk_combinations.py b/..._computation/test_jax_ntk_combinations.py → ...sts/analysis/test_jax_ntk_combinations.py
@@ -32,8 +32,8 @@
 from jax import random
 from papyrus.utils.matrix_utils import flatten_rank_4_tensor
 
+from znnl.analysis import JAXNTKCombinations
 from znnl.models import FlaxModel
-from znnl.ntk_computation import JAXNTKCombinations
 
 
 class FlaxTestModule(nn.Module):

diff --git a/...k_computation/test_jax_ntk_subsampling.py → ...ests/analysis/test_jax_ntk_subsampling.py b/...k_computation/test_jax_ntk_subsampling.py → ...ests/analysis/test_jax_ntk_subsampling.py
@@ -30,8 +30,8 @@
 from flax import linen as nn
 from jax import random
 
+from znnl.analysis import JAXNTKSubsampling
 from znnl.models import FlaxModel
-from znnl.ntk_computation import JAXNTKSubsampling
 
 
 class FlaxTestModule(nn.Module):

diff --git a/CI/unit_tests/models/_test_huggingface_flax_model.py b/CI/unit_tests/models/_test_huggingface_flax_model.py
@@ -29,8 +29,8 @@
 from jax import random
 from transformers import FlaxResNetForImageClassification, ResNetConfig
 
+from znnl.analysis import JAXNTKComputation
 from znnl.models import HuggingFaceFlaxModel
-from znnl.ntk_computation import JAXNTKComputation
 
 
 class TestFlaxHFModule:

diff --git a/CI/unit_tests/models/test_flax_model.py b/CI/unit_tests/models/test_flax_model.py
@@ -34,8 +34,8 @@
 from flax import linen as nn
 from jax import random
 
+from znnl.analysis import JAXNTKComputation
 from znnl.models import FlaxModel
-from znnl.ntk_computation import JAXNTKComputation
 
 
 class FlaxTestModule(nn.Module):

diff --git a/CI/unit_tests/models/test_nt_model.py b/CI/unit_tests/models/test_nt_model.py
@@ -34,8 +34,8 @@
 from jax import random
 from neural_tangents import stax
 
+from znnl.analysis import JAXNTKComputation
 from znnl.models import NTModel
-from znnl.ntk_computation import JAXNTKComputation
 
 
 class TestNTModule:

diff --git a/CI/unit_tests/optimizers/test_trace_optimizer.py b/CI/unit_tests/optimizers/test_trace_optimizer.py
@@ -32,9 +32,9 @@
 import jax.numpy as np
 from neural_tangents import stax
 
+from znnl.analysis import JAXNTKComputation
 from znnl.data import MNISTGenerator
 from znnl.models import NTModel
-from znnl.ntk_computation import JAXNTKComputation
 from znnl.optimizers import TraceOptimizer
 
 

diff --git a/CI/unit_tests/training_recording/test_jax_recorder.py b/CI/unit_tests/training_recording/test_jax_recorder.py
@@ -33,8 +33,8 @@
 from numpy.testing import assert_raises
 from papyrus.measurements import Accuracy, Loss, NTKTrace
 
+from znnl.analysis import JAXNTKComputation
 from znnl.models import FlaxModel
-from znnl.ntk_computation import JAXNTKComputation
 from znnl.training_recording import JaxRecorder
 
 

diff --git a/examples/Computing-Collective-Variables.ipynb b/examples/Computing-Collective-Variables.ipynb
@@ -165,7 +165,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ntk_computation = nl.ntk_computation.JAXNTKComputation(\n",
+    "ntk_computation = nl.analysis.JAXNTKComputation(\n",
     "    apply_fn=fuel_model.ntk_apply_fn, \n",
     "    batch_size=314,\n",
     ")\n",

diff --git a/examples/Contrastive-Loss.ipynb b/examples/Contrastive-Loss.ipynb
@@ -192,7 +192,7 @@
     "    update_rate=1, \n",
     "    chunk_size=1e10 # Big Chunk-size to prevent saving the recordings.\n",
     ")\n",
-    "ntk_computation = znnl.ntk_computation.JAXNTKComputation(\n",
+    "ntk_computation = znnl.analysis.JAXNTKComputation(\n",
     "    apply_fn=model.ntk_apply_fn, \n",
     "    batch_size=10,\n",
     ")\n",
@@ -461,7 +461,7 @@
     "    update_rate=1, \n",
     "    chunk_size=1e10 # Big Chunk-size to prevent saving the recordings.\n",
     ")\n",
-    "ntk_computation = znnl.ntk_computation.JAXNTKComputation(\n",
+    "ntk_computation = znnl.analysis.JAXNTKComputation(\n",
     "    apply_fn=model.ntk_apply_fn, \n",
     "    batch_size=10,\n",
     ")\n",

diff --git a/examples/ResNet-Example.ipynb b/examples/ResNet-Example.ipynb
@@ -174,7 +174,7 @@
     "    ],\n",
     "    update_rate=1, \n",
     ")\n",
-    "ntk_computation = nl.ntk_computation.JAXNTKComputation(\n",
+    "ntk_computation = nl.analysis.JAXNTKComputation(\n",
     "    apply_fn=model.ntk_apply_fn, \n",
     "    batch_size=10,    \n",
     ")\n",

diff --git a/znnl/agents/approximate_maximum_entropy.py b/znnl/agents/approximate_maximum_entropy.py
@@ -30,9 +30,9 @@
 
 from znnl.agents.agent import Agent
 from znnl.analysis.entropy import EntropyAnalysis
+from znnl.analysis.jax_ntk import JAXNTKComputation
 from znnl.data import DataGenerator
 from znnl.models import JaxModel
-from znnl.ntk_computation.jax_ntk import JAXNTKComputation
 from znnl.utils.prng import PRNGKey
 
 

diff --git a/znnl/analysis/__init__.py b/znnl/analysis/__init__.py
@@ -27,10 +27,18 @@
 
 from znnl.analysis.eigensystem import EigenSpaceAnalysis
 from znnl.analysis.entropy import EntropyAnalysis
+from znnl.analysis.jax_ntk import JAXNTKComputation
+from znnl.analysis.jax_ntk_classwise import JAXNTKClassWise
+from znnl.analysis.jax_ntk_combinations import JAXNTKCombinations
+from znnl.analysis.jax_ntk_subsampling import JAXNTKSubsampling
 from znnl.analysis.loss_fn_derivative import LossDerivative
 
 __all__ = [
     EntropyAnalysis.__name__,
     EigenSpaceAnalysis.__name__,
     LossDerivative.__name__,
+    JAXNTKComputation.__name__,
+    JAXNTKClassWise.__name__,
+    JAXNTKSubsampling.__name__,
+    JAXNTKCombinations.__name__,
 ]
diff --git a/znnl/ntk_computation/jax_ntk.py → znnl/analysis/jax_ntk.py b/znnl/ntk_computation/jax_ntk.py → znnl/analysis/jax_ntk.py
diff --git a/znnl/ntk_computation/jax_ntk_classwise.py → znnl/analysis/jax_ntk_classwise.py b/znnl/ntk_computation/jax_ntk_classwise.py → znnl/analysis/jax_ntk_classwise.py
@@ -32,7 +32,7 @@
 from jax import random, vmap
 from jax.tree_util import tree_map as jmap
 
-from znnl.ntk_computation.jax_ntk import JAXNTKComputation
+from znnl.analysis.jax_ntk import JAXNTKComputation
 
 
 class JAXNTKClassWise(JAXNTKComputation):

diff --git a/znnl/ntk_computation/jax_ntk_combinations.py → znnl/analysis/jax_ntk_combinations.py b/znnl/ntk_computation/jax_ntk_combinations.py → znnl/analysis/jax_ntk_combinations.py
@@ -32,7 +32,7 @@
 import neural_tangents as nt
 from papyrus.utils.matrix_utils import flatten_rank_4_tensor, unflatten_rank_4_tensor
 
-from znnl.ntk_computation.jax_ntk import JAXNTKComputation
+from znnl.analysis.jax_ntk import JAXNTKComputation
 
 
 class JAXNTKCombinations(JAXNTKComputation):

diff --git a/znnl/ntk_computation/jax_ntk_subsampling.py → znnl/analysis/jax_ntk_subsampling.py b/znnl/ntk_computation/jax_ntk_subsampling.py → znnl/analysis/jax_ntk_subsampling.py
@@ -32,7 +32,7 @@
 from jax import random
 from jax.tree_util import tree_map as jmap
 
-from znnl.ntk_computation.jax_ntk import JAXNTKComputation
+from znnl.analysis.jax_ntk import JAXNTKComputation
 
 
 class JAXNTKSubsampling(JAXNTKComputation):

diff --git a/znnl/ntk_computation/__init__.py b/znnl/ntk_computation/__init__.py
diff --git a/znnl/training_recording/jax_recording.py b/znnl/training_recording/jax_recording.py
@@ -37,10 +37,10 @@
 from znnl.accuracy_functions.accuracy_function import AccuracyFunction
 from znnl.analysis.eigensystem import EigenSpaceAnalysis
 from znnl.analysis.entropy import EntropyAnalysis
+from znnl.analysis.jax_ntk import JAXNTKComputation
 from znnl.analysis.loss_fn_derivative import LossDerivative
 from znnl.loss_functions import SimpleLoss
 from znnl.models.jax_model import JaxModel
-from znnl.ntk_computation.jax_ntk import JAXNTKComputation
 from znnl.training_recording.data_storage import DataStorage
 from znnl.utils.matrix_utils import (
     calculate_trace,

diff --git a/znnl/training_recording/papyrus_jax_recording.py b/znnl/training_recording/papyrus_jax_recording.py
@@ -31,8 +31,8 @@
 from papyrus.measurements import BaseMeasurement
 from papyrus.recorders import BaseRecorder
 
+from znnl.analysis import JAXNTKComputation
 from znnl.models import JaxModel
-from znnl.ntk_computation import JAXNTKComputation
 
 
 class JaxRecorder(BaseRecorder):

diff --git a/znnl/training_strategies/loss_aware_reservoir.py b/znnl/training_strategies/loss_aware_reservoir.py
@@ -35,9 +35,9 @@
 from tqdm import trange
 
 from znnl.accuracy_functions.accuracy_function import AccuracyFunction
+from znnl.analysis.jax_ntk import JAXNTKComputation
 from znnl.distance_metrics import DistanceMetric
 from znnl.models.jax_model import JaxModel
-from znnl.ntk_computation.jax_ntk import JAXNTKComputation
 from znnl.optimizers.trace_optimizer import TraceOptimizer
 from znnl.training_recording import JaxRecorder
 from znnl.training_strategies.recursive_mode import RecursiveMode

diff --git a/znnl/training_strategies/partitioned_training.py b/znnl/training_strategies/partitioned_training.py
@@ -32,8 +32,8 @@
 from tqdm import trange
 
 from znnl.accuracy_functions.accuracy_function import AccuracyFunction
+from znnl.analysis.jax_ntk import JAXNTKComputation
 from znnl.models.jax_model import JaxModel
-from znnl.ntk_computation.jax_ntk import JAXNTKComputation
 from znnl.optimizers.trace_optimizer import TraceOptimizer
 from znnl.training_recording import JaxRecorder
 from znnl.training_strategies.recursive_mode import RecursiveMode

diff --git a/znnl/training_strategies/simple_training.py b/znnl/training_strategies/simple_training.py
@@ -36,8 +36,8 @@
 from tqdm import trange
 
 from znnl.accuracy_functions.accuracy_function import AccuracyFunction
+from znnl.analysis.jax_ntk import JAXNTKComputation
 from znnl.models.jax_model import JaxModel
-from znnl.ntk_computation.jax_ntk import JAXNTKComputation
 from znnl.optimizers.trace_optimizer import TraceOptimizer
 from znnl.training_recording import JaxRecorder
 from znnl.training_strategies.recursive_mode import RecursiveMode