Reduce the peak memory usage of skillmodels by hmgaudecker · Pull Request #76 · OpenSourceEconomics/skillmodels

Change citation year so we can open PR.

hmgaudecker

committed

Add a version of mem consumption test that fails if there are changes in the repo.

hmgaudecker

committed

Add table tracking memory usage.

hmgaudecker

committed

Add table tracking memory usage.

hmgaudecker

committed

Remove jitting altogether, tiny refactorings.

hmgaudecker

committed

Jit _scan_body

hmgaudecker

committed

Remove jitting again, do not assign new object to by just keeping the second element of the tuple returned by jax.lax.scan.

hmgaudecker

committed

Only pass around the transition function instead of the info dict where possible.

hmgaudecker

committed

jit scan_body.

hmgaudecker

committed

Update MEMTRACKER, still opposite effect.

hmgaudecker

committed

Back to start.

hmgaudecker

committed

Decorate kalman_update with jax.remat.

hmgaudecker

committed

Update memtracker.

hmgaudecker

committed

Decorate _calculate_sigma_points with jax.remat.

hmgaudecker

committed

Use checkpoint() alias to remat, can use due to being inside lax.scan, as per https://jax.readthedocs.io/en/latest/_autosummary/jax.checkpoint.html\#jax.checkpoint.

hmgaudecker

committed

Use checkpoint() on kalman_predict as well, requires changing to be its first argument and using in tests.

hmgaudecker

committed

Update MEMTRACKER

hmgaudecker

committed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Reduce the peak memory usage of skillmodels#76

Reduce the peak memory usage of skillmodels#76
hmgaudecker merged 23 commits intomainfrom
reduce-mem-usage

Commits on Sep 10, 2024

Change citation year so we can open PR.

Add a version of mem consumption test that fails if there are changes in the repo.

Add table tracking memory usage.

Add table tracking memory usage.

Remove jitting altogether, tiny refactorings.

Jit _scan_body

Remove jitting again, do not assign new object to by just keeping the second element of the tuple returned by jax.lax.scan.

Only pass around the transition function instead of the info dict where possible.

jit scan_body.

Update MEMTRACKER, still opposite effect.

Back to start.

Decorate kalman_update with jax.remat.

Update memtracker.

Decorate _calculate_sigma_points with jax.remat.

Use checkpoint() alias to remat, can use due to being inside lax.scan, as per https://jax.readthedocs.io/en/latest/_autosummary/jax.checkpoint.html\#jax.checkpoint.

Use checkpoint() on kalman_predict as well, requires changing to be its first argument and using in tests.

Update MEMTRACKER

Commits on Sep 11, 2024

Provision for testing on GPUs, timing not useful yet.

Fix environment and actions.

Add pixi.lock

Merge branch 'main' into reduce-mem-usage

Refactor likelihood module (#77)

Remove file used for tracking progress.