diff --git a/.idea/diffrax_STLA.iml b/.idea/diffrax_STLA.iml
deleted file mode 100644
index aad402c4..00000000
--- a/.idea/diffrax_STLA.iml
+++ /dev/null
@@ -1,10 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<module version="4">
-  <component name="PyDocumentationSettings">
-    <option name="format" value="PLAIN" />
-    <option name="myDocStringFormat" value="Plain" />
-  </component>
-  <component name="TestRunnerService">
-    <option name="PROJECT_TEST_RUNNER" value="py.test" />
-  </component>
-</module>
\ No newline at end of file
diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml
deleted file mode 100644
index 105ce2da..00000000
--- a/.idea/inspectionProfiles/profiles_settings.xml
+++ /dev/null
@@ -1,6 +0,0 @@
-<component name="InspectionProjectProfileManager">
-  <settings>
-    <option name="USE_PROJECT_PROFILE" value="false" />
-    <version value="1.0" />
-  </settings>
-</component>
\ No newline at end of file
diff --git a/.idea/misc.xml b/.idea/misc.xml
deleted file mode 100644
index 4c223141..00000000
--- a/.idea/misc.xml
+++ /dev/null
@@ -1,7 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<project version="4">
-  <component name="Black">
-    <option name="sdkName" value="Python 3.11 (diffrax_STLA)" />
-  </component>
-  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.11 (diffrax_STLA)" project-jdk-type="Python SDK" />
-</project>
\ No newline at end of file
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
deleted file mode 100644
index 35eb1ddf..00000000
--- a/.idea/vcs.xml
+++ /dev/null
@@ -1,6 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<project version="4">
-  <component name="VcsDirectoryMappings">
-    <mapping directory="" vcs="Git" />
-  </component>
-</project>
\ No newline at end of file
diff --git a/.idea/workspace.xml b/.idea/workspace.xml
deleted file mode 100644
index 9666f9a8..00000000
--- a/.idea/workspace.xml
+++ /dev/null
@@ -1,262 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<project version="4">
-  <component name="AutoImportSettings">
-    <option name="autoReloadType" value="SELECTIVE" />
-  </component>
-  <component name="ChangeListManager">
-    <list default="true" id="14f04cc7-7857-4a3e-8787-78e632edfd0f" name="Changes" comment="">
-      <change beforePath="$PROJECT_DIR$/test/test_align.py" beforeDir="false" afterPath="$PROJECT_DIR$/test/test_align.py" afterDir="false" />
-    </list>
-    <option name="SHOW_DIALOG" value="false" />
-    <option name="HIGHLIGHT_CONFLICTS" value="true" />
-    <option name="HIGHLIGHT_NON_ACTIVE_CHANGELIST" value="false" />
-    <option name="LAST_RESOLUTION" value="IGNORE" />
-  </component>
-  <component name="FileTemplateManagerImpl">
-    <option name="RECENT_TEMPLATES">
-      <list>
-        <option value="Python Script" />
-        <option value="Jupyter Notebook" />
-      </list>
-    </option>
-  </component>
-  <component name="Git.Settings">
-    <option name="RECENT_GIT_ROOT_PATH" value="$PROJECT_DIR$" />
-  </component>
-  <component name="HighlightingSettingsPerFile">
-    <setting file="file://$PROJECT_DIR$/../JAX_trial/venv2/lib/python3.10/site-packages/equinox/internal/_omega.py" root0="SKIP_INSPECTION" />
-  </component>
-  <component name="MarkdownSettingsMigration">
-    <option name="stateVersion" value="1" />
-  </component>
-  <component name="ProjectColorInfo">{
-  &quot;associatedIndex&quot;: 8
-}</component>
-  <component name="ProjectId" id="2WZ10GV49U8b5jdhCvqKaLBYqle" />
-  <component name="ProjectViewState">
-    <option name="hideEmptyMiddlePackages" value="true" />
-    <option name="showLibraryContents" value="true" />
-  </component>
-  <component name="PropertiesComponent">{
-  &quot;keyToString&quot;: {
-    &quot;RunOnceActivity.OpenProjectViewOnStart&quot;: &quot;true&quot;,
-    &quot;RunOnceActivity.ShowReadmeOnStart&quot;: &quot;true&quot;,
-    &quot;WebServerToolWindowFactoryState&quot;: &quot;false&quot;,
-    &quot;last_opened_file_path&quot;: &quot;/home/andy/PycharmProjects/diffrax_STLA/diffrax/custom_types.py&quot;,
-    &quot;node.js.detected.package.eslint&quot;: &quot;true&quot;,
-    &quot;node.js.detected.package.tslint&quot;: &quot;true&quot;,
-    &quot;node.js.selected.package.eslint&quot;: &quot;(autodetect)&quot;,
-    &quot;node.js.selected.package.tslint&quot;: &quot;(autodetect)&quot;,
-    &quot;nodejs_package_manager_path&quot;: &quot;npm&quot;,
-    &quot;settings.editor.selected.configurable&quot;: &quot;com.jetbrains.python.configuration.PyActiveSdkModuleConfigurable&quot;,
-    &quot;vue.rearranger.settings.migration&quot;: &quot;true&quot;
-  },
-  &quot;keyToStringList&quot;: {
-    &quot;com.intellij.ide.scratch.ScratchImplUtil$2/New Scratch File&quot;: [
-      &quot;PLAIN_TEXT&quot;
-    ]
-  }
-}</component>
-  <component name="RecentsManager">
-    <key name="CopyFile.RECENT_KEYS">
-      <recent name="$PROJECT_DIR$/test" />
-      <recent name="$PROJECT_DIR$/diffrax/STLA" />
-    </key>
-    <key name="MoveFile.RECENT_KEYS">
-      <recent name="$PROJECT_DIR$/diffrax/brownian" />
-    </key>
-  </component>
-  <component name="RunManager" selected="Python tests.pytest for test_align.test_shape">
-    <configuration name="pytest for test_align.test_convergence" type="tests" factoryName="py.test" nameIsGenerated="true">
-      <module name="diffrax_STLA" />
-      <option name="INTERPRETER_OPTIONS" value="" />
-      <option name="PARENT_ENVS" value="true" />
-      <option name="SDK_HOME" value="" />
-      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/test" />
-      <option name="IS_MODULE_SDK" value="true" />
-      <option name="ADD_CONTENT_ROOTS" value="true" />
-      <option name="ADD_SOURCE_ROOTS" value="true" />
-      <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
-      <option name="_new_keywords" value="&quot;&quot;" />
-      <option name="_new_parameters" value="&quot;&quot;" />
-      <option name="_new_additionalArguments" value="&quot;&quot;" />
-      <option name="_new_target" value="&quot;test_align.test_convergence&quot;" />
-      <option name="_new_targetType" value="&quot;PYTHON&quot;" />
-      <method v="2" />
-    </configuration>
-    <configuration name="pytest for test_align.test_shape" type="tests" factoryName="py.test" temporary="true" nameIsGenerated="true">
-      <module name="diffrax_STLA" />
-      <option name="INTERPRETER_OPTIONS" value="" />
-      <option name="PARENT_ENVS" value="true" />
-      <option name="SDK_HOME" value="" />
-      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/test" />
-      <option name="IS_MODULE_SDK" value="true" />
-      <option name="ADD_CONTENT_ROOTS" value="true" />
-      <option name="ADD_SOURCE_ROOTS" value="true" />
-      <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
-      <option name="_new_keywords" value="&quot;&quot;" />
-      <option name="_new_parameters" value="&quot;&quot;" />
-      <option name="_new_additionalArguments" value="&quot;&quot;" />
-      <option name="_new_target" value="&quot;test_align.test_shape&quot;" />
-      <option name="_new_targetType" value="&quot;PYTHON&quot;" />
-      <method v="2" />
-    </configuration>
-    <configuration name="pytest for test_stla.test_conditional_statistics" type="tests" factoryName="py.test" temporary="true" nameIsGenerated="true">
-      <module name="diffrax_STLA" />
-      <option name="INTERPRETER_OPTIONS" value="" />
-      <option name="PARENT_ENVS" value="true" />
-      <option name="SDK_HOME" value="" />
-      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/test" />
-      <option name="IS_MODULE_SDK" value="true" />
-      <option name="ADD_CONTENT_ROOTS" value="true" />
-      <option name="ADD_SOURCE_ROOTS" value="true" />
-      <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
-      <option name="_new_keywords" value="&quot;&quot;" />
-      <option name="_new_parameters" value="&quot;&quot;" />
-      <option name="_new_additionalArguments" value="&quot;&quot;" />
-      <option name="_new_target" value="&quot;test_stla.test_conditional_statistics&quot;" />
-      <option name="_new_targetType" value="&quot;PYTHON&quot;" />
-      <method v="2" />
-    </configuration>
-    <configuration name="pytest for test_stla.test_shape_and_dtype" type="tests" factoryName="py.test" temporary="true" nameIsGenerated="true">
-      <module name="diffrax_STLA" />
-      <option name="INTERPRETER_OPTIONS" value="" />
-      <option name="PARENT_ENVS" value="true" />
-      <option name="SDK_HOME" value="" />
-      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/test" />
-      <option name="IS_MODULE_SDK" value="true" />
-      <option name="ADD_CONTENT_ROOTS" value="true" />
-      <option name="ADD_SOURCE_ROOTS" value="true" />
-      <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
-      <option name="_new_keywords" value="&quot;&quot;" />
-      <option name="_new_parameters" value="&quot;&quot;" />
-      <option name="_new_additionalArguments" value="&quot;&quot;" />
-      <option name="_new_target" value="&quot;test_stla.test_shape_and_dtype&quot;" />
-      <option name="_new_targetType" value="&quot;PYTHON&quot;" />
-      <method v="2" />
-    </configuration>
-    <configuration name="pytest for test.test_detest.test_a" type="tests" factoryName="py.test" temporary="true" nameIsGenerated="true">
-      <module name="diffrax_STLA" />
-      <option name="INTERPRETER_OPTIONS" value="" />
-      <option name="PARENT_ENVS" value="true" />
-      <option name="SDK_HOME" value="" />
-      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$" />
-      <option name="IS_MODULE_SDK" value="true" />
-      <option name="ADD_CONTENT_ROOTS" value="true" />
-      <option name="ADD_SOURCE_ROOTS" value="true" />
-      <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
-      <option name="_new_keywords" value="&quot;&quot;" />
-      <option name="_new_parameters" value="&quot;&quot;" />
-      <option name="_new_additionalArguments" value="&quot;&quot;" />
-      <option name="_new_target" value="&quot;test.test_detest.test_a&quot;" />
-      <option name="_new_targetType" value="&quot;PYTHON&quot;" />
-      <method v="2" />
-    </configuration>
-    <configuration name="pytest for test.test_detest.test_b" type="tests" factoryName="py.test" temporary="true" nameIsGenerated="true">
-      <module name="diffrax_STLA" />
-      <option name="INTERPRETER_OPTIONS" value="" />
-      <option name="PARENT_ENVS" value="true" />
-      <option name="SDK_HOME" value="" />
-      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$" />
-      <option name="IS_MODULE_SDK" value="true" />
-      <option name="ADD_CONTENT_ROOTS" value="true" />
-      <option name="ADD_SOURCE_ROOTS" value="true" />
-      <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
-      <option name="_new_keywords" value="&quot;&quot;" />
-      <option name="_new_parameters" value="&quot;&quot;" />
-      <option name="_new_additionalArguments" value="&quot;&quot;" />
-      <option name="_new_target" value="&quot;test.test_detest.test_b&quot;" />
-      <option name="_new_targetType" value="&quot;PYTHON&quot;" />
-      <method v="2" />
-    </configuration>
-    <recent_temporary>
-      <list>
-        <item itemvalue="Python tests.pytest for test_align.test_shape" />
-        <item itemvalue="Python tests.pytest for test.test_detest.test_a" />
-        <item itemvalue="Python tests.pytest for test.test_detest.test_b" />
-        <item itemvalue="Python tests.pytest for test_stla.test_shape_and_dtype" />
-        <item itemvalue="Python tests.pytest for test_stla.test_conditional_statistics" />
-      </list>
-    </recent_temporary>
-  </component>
-  <component name="SpellCheckerSettings" RuntimeDictionaries="0" Folders="0" CustomDictionaries="0" DefaultDictionary="application-level" UseSingleDictionary="true" transferred="true" />
-  <component name="TaskManager">
-    <task active="true" id="Default" summary="Default task">
-      <changelist id="14f04cc7-7857-4a3e-8787-78e632edfd0f" name="Changes" comment="" />
-      <created>1696925063311</created>
-      <option name="number" value="Default" />
-      <option name="presentableId" value="Default" />
-      <updated>1696925063311</updated>
-      <workItem from="1696925064528" duration="691000" />
-      <workItem from="1696943975882" duration="9502000" />
-      <workItem from="1697019632032" duration="3082000" />
-      <workItem from="1697033118441" duration="660000" />
-      <workItem from="1697049484735" duration="3001000" />
-      <workItem from="1697105998587" duration="4407000" />
-      <workItem from="1697187166227" duration="7687000" />
-      <workItem from="1697216500461" duration="3577000" />
-      <workItem from="1697298976912" duration="12260000" />
-      <workItem from="1697308246606" duration="1423000" />
-      <workItem from="1697319089411" duration="422000" />
-      <workItem from="1697319526632" duration="128000" />
-      <workItem from="1697447318106" duration="11975000" />
-      <workItem from="1697470622602" duration="3131000" />
-      <workItem from="1697532110281" duration="20390000" />
-      <workItem from="1697620424878" duration="25127000" />
-      <workItem from="1697661644085" duration="2784000" />
-      <workItem from="1697709913824" duration="2152000" />
-      <workItem from="1697723805535" duration="7749000" />
-      <workItem from="1697733881249" duration="4439000" />
-      <workItem from="1697744698279" duration="1368000" />
-      <workItem from="1697809882763" duration="11042000" />
-      <workItem from="1697883424224" duration="20799000" />
-      <workItem from="1698059173515" duration="8341000" />
-      <workItem from="1698143458834" duration="14734000" />
-      <workItem from="1698227373267" duration="5320000" />
-      <workItem from="1698236648679" duration="7444000" />
-      <workItem from="1698247738314" duration="2271000" />
-      <workItem from="1698259117206" duration="1455000" />
-      <workItem from="1698262059467" duration="412000" />
-      <workItem from="1698262829493" duration="3302000" />
-      <workItem from="1698312687839" duration="14466000" />
-      <workItem from="1698398668683" duration="3048000" />
-      <workItem from="1698409199887" duration="1606000" />
-      <workItem from="1698410978363" duration="2160000" />
-      <workItem from="1698415856292" duration="160000" />
-      <workItem from="1698428645925" duration="3605000" />
-      <workItem from="1698660571974" duration="4853000" />
-      <workItem from="1698688373126" duration="3695000" />
-      <workItem from="1698744641774" duration="6622000" />
-      <workItem from="1698751357985" duration="14910000" />
-      <workItem from="1698835148582" duration="1298000" />
-      <workItem from="1698847705745" duration="4256000" />
-      <workItem from="1698861272181" duration="1708000" />
-      <workItem from="1698866123705" duration="2693000" />
-      <workItem from="1698869199580" duration="2104000" />
-      <workItem from="1698873861733" duration="4921000" />
-      <workItem from="1698922202647" duration="1646000" />
-      <workItem from="1698923906663" duration="2864000" />
-      <workItem from="1699093251958" duration="505000" />
-      <workItem from="1699372345481" duration="3956000" />
-      <workItem from="1699376754915" duration="1209000" />
-    </task>
-    <servers />
-  </component>
-  <component name="TypeScriptGeneratedFilesManager">
-    <option name="version" value="3" />
-  </component>
-  <component name="com.intellij.coverage.CoverageDataManagerImpl">
-    <SUITE FILE_PATH="coverage/diffrax_STLA$pytest_for_test_brownian_test_statistics.coverage" NAME="pytest for test_brownian.test_statistics Coverage Results" MODIFIED="1698855084730" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/test" />
-    <SUITE FILE_PATH="coverage/diffrax_STLA$pytest_for_test_stla_test_statistics.coverage" NAME="pytest for test_stla.test_statistics Coverage Results" MODIFIED="1698861445853" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/test" />
-    <SUITE FILE_PATH="coverage/diffrax_STLA$pytest_for_test_align_test_convergence.coverage" NAME="pytest for test_align.test_convergence Coverage Results" MODIFIED="1699375899419" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/test" />
-    <SUITE FILE_PATH="coverage/diffrax_STLA$pytest_for_test_stla_test_shape_and_dtype.coverage" NAME="pytest for test_stla.test_shape_and_dtype Coverage Results" MODIFIED="1698862851423" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/test" />
-    <SUITE FILE_PATH="coverage/diffrax_STLA$pytest_for_test_align_test_shape_dtype.coverage" NAME="pytest for test_align.test_shape_dtype Coverage Results" MODIFIED="1698777302719" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/test" />
-    <SUITE FILE_PATH="coverage/diffrax_STLA$pytest_for_test_brownian_test_shape_and_dtype.coverage" NAME="pytest for test_brownian.test_shape_and_dtype Coverage Results" MODIFIED="1698854950633" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/test" />
-    <SUITE FILE_PATH="coverage/diffrax_STLA$pytest_for_test_test_detest_test_a.coverage" NAME="pytest for test.test_detest.test_a Coverage Results" MODIFIED="1698870800000" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$" />
-    <SUITE FILE_PATH="coverage/diffrax_STLA$pytest_for_test_test_adjoint_test_against.coverage" NAME="pytest for test.test_adjoint.test_against Coverage Results" MODIFIED="1698772328998" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$" />
-    <SUITE FILE_PATH="coverage/diffrax_STLA$pytest_for_test_brownian_test_conditional_statistics.coverage" NAME="pytest for test_brownian.test_conditional_statistics Coverage Results" MODIFIED="1698855094085" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/test" />
-    <SUITE FILE_PATH="coverage/diffrax_STLA$pytest_for_test_stla_test_conditional_statistics.coverage" NAME="pytest for test_stla.test_conditional_statistics Coverage Results" MODIFIED="1698861766847" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/test" />
-    <SUITE FILE_PATH="coverage/diffrax_STLA$pytest_for_test_test_detest_test_b.coverage" NAME="pytest for test.test_detest.test_b Coverage Results" MODIFIED="1698869072989" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$" />
-    <SUITE FILE_PATH="coverage/diffrax_STLA$pytest_for_test_align_test_shape.coverage" NAME="pytest for test_align.test_shape Coverage Results" MODIFIED="1699376770181" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="true" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/test" />
-  </component>
-</project>
\ No newline at end of file