Skip to content

Flax/nnx backend #440

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 212 commits into from
Aug 16, 2024
Merged

Flax/nnx backend #440

merged 212 commits into from
Aug 16, 2024

Conversation

frazane
Copy link
Contributor

@frazane frazane commented Mar 6, 2024 •
edited by thomaspinder
Loading

Description

Introduces flax / nnx as backend for the library.

Tasks

  • pass all standard tests
  • pass integration tests
  • check and eventually update all docstrings
  • pass docstrings tests
  • run all documentation notebooks
  • adapt markdown cells in documentation notebooks
  • new docs pages

Running documentation notebooks

  • barycentres
  • bayesian_optimisation
  • classification
  • collapsed_vi
  • constructing_new_kernels
  • decision_making
  • deep_kernels
  • graph_kernels
  • intro_to_gps
  • intro_to_kernels
  • likelihoods_guide
  • oceanmodelling
  • poisson
  • regression
  • uncollapsed_vi
  • utils
  • yacht

frazane added 30 commits January 6, 2024 13:24
@frazane
add flax v0.8.0 to deps, temporarily from github main branch
f7ff6b1
@frazane
main gps objects as nnx modules
04ce459
@frazane
integrators as nnx dataclasses and some static typing refactoring
2dc9479
@frazane
likelihoods as nnx dataclasses modules and some static typing refacto…
7efbd20 
…ring
@frazane
small refactoring
ae034e2
@frazane
mean functions as nnx dataclasses modules and some refactoring
c530481
@frazane
bugfix
7d14eaf
@frazane
objectives as nnx dataclasses modules
40891f7
@frazane
variational families with nnx
e749bc6
@frazane
kernels base with nnx
3568d9c
@frazane
wip stationary kernels
28c3b9e
@frazane
wip nonstationary kernels
49cc6e0
@frazane
wip non euclidean kernels
9d2567e
@frazane
computations with nnx
f18701c
@frazane
rff with nnx
61bf22b
@frazane
bugfix
f0736f7
@frazane
Merge branch 'main' into flax-backend
64f4f3b
@frazane
stationary kernels as normal classes
6038458
@frazane
nonstationary kernels as normal classes
447dad1
@frazane
noneuclidean kernels as normal classes
2d343c2
@frazane
rff as standard class + stationary kernel abstract class for static t…
af33559 
…yping
@frazane
started work on parameters
cfabc4c
@frazane
more objects as normal classes
645c111
@frazane
gps as normal classes
c536b1a
@frazane
integrators as normal classes
cac712e
@frazane
dataset is not a pytree
f50ffe7
@frazane
removed superfluous inits
7e1a23e
@frazane
register dataset as pytree
5d21926
@frazane
use parameters here and there
8133850
@frazane
set active_dims default to 1
4333f0b
thomaspinder and others added 26 commits August 15, 2024 22:29
@thomaspinder
bump deps
939f903
@thomaspinder
Update poetry
8ddee1f
@thomaspinder
Update poetry
767cef2
@thomaspinder
Merge branch 'fix-doc-build' into backend_doc
46be074 
Signed-off-by: Thomas Pinder <tompinder@live.co.uk>
@thomaspinder
Fix shutil
81a6b98
@thomaspinder
Drop flax base
e38d5bb
@thomaspinder
Merge pull request #461 from JaxGaussianProcesses/fix-doc-build
0d21a31 
Update docs build process
@thomaspinder
Rebase
382098c
@thomaspinder
Merge branch 'backend_doc' of github.com:JaxGaussianProcesses/GPJax i…
a5bca61 
…nto backend_doc
@thomaspinder
Merge pull request #462 from JaxGaussianProcesses/backend_doc
4383d1d 
Backend doc
@frazane
add scikit-learn dependency for docs
0c073ac
@frazane
bugfix: change directory before running jupytext
91022b9
@frazane
use local mpl style file
7d73c27
@frazane
do not use MCMC for classification (it is *very* slow)
c07ead4
@frazane
[skip-ci] update github workflows for docs
4ebc673
@thomaspinder
Fix split
6450603
@thomaspinder
Fix split
4ee615e
@thomaspinder
Fix split
2d323b2
@thomaspinder
Fix xdoctest
5b9f058
@thomaspinder
Fix doc
6fcb83d
@thomaspinder
Add serial build
02bfd01
@thomaspinder
Update parameters transform and backend doc
7e628f0
@thomaspinder
Update parameters transform and backend doc
172e61e
@thomaspinder
Merge pull request #465 from JaxGaussianProcesses/fix-doc-build
ab86ac5 
Fix doc build
@thomaspinder
Bump Python
cc62521
@thomaspinder
Merge branch 'flax-backend-normalclasses' of github.com:JaxGaussianPr…
20953c6 
…ocesses/GPJax into flax-backend-normalclasses
@thomaspinder thomaspinder changed the title [WIP] Flax/nnx backend Flax/nnx backend Aug 16, 2024
@thomaspinder thomaspinder merged commit 9ba68a4 into main Aug 16, 2024
13 checks passed
@thomaspinder thomaspinder deleted the flax-backend-normalclasses branch August 16, 2024 19:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@thomaspinder thomaspinder thomaspinder approved these changes

Assignees
No one assigned
Labels
breaking Breaking Changes enhancement New feature or request work-in-progress
Projects
None yet
Milestone
0.9.0
Development

Successfully merging this pull request may close these issues.
4 participants
@frazane @thomaspinder @daniel-dodd @Thomas-Christie