umangjpatel / kerax Goto Github PK

Description
Add regularization layers such as Inverted Dropout and Batch Norm

Description
Implement JAX vmap structure for batching

Solution
Use JAX vmap for batched gradient descent implementation.

Additional context
Refer to https://jax.readthedocs.io/en/latest/index.html

Description
Implement various weight initialization schemes

Solution
Use jax.nn module for implementation notes.

Description
PyTorch support inefficient. Takes too much time.

Describe the solution you'd like
Reimplement the library in JAX

Additional context
Refer to https://github.com/google/jax

Description
Update README and add experiments.

Solution
Create experiments or examples package to work

Description
Constant accuracy and large losses while using CNN APIs

Expected behaviour
Accuracy should increase while loss should decrease

Description
Softmax classification algorithm

Description
Add optimization algorithms

Solution
Use JAX experimental module

Description of feature request
Multiclass / Multinomial classification problems not supported.

Solution proposed
Implement softmax function

Additional context
Refer to docs.fast.ai, PyTorch and Keras documentation for implementation details

Description
Implement a databunch-like API for loading image data.

Solution
Use PyTorch and TensorFlow data loading libraries

Additional context
Refer to jax.readthedocs.io for more info.

Description
JAX imports displays NotImplementedError

Expected behavior
To use JAX-numpy in accordance with the original NumPy package.

Proposed Solution
Replace JAX computation engine with PyTorch package.

Is your feature request related to a problem? Please describe.
After training large models, we shall be able to save and reload them for further experiments.

Describe the solution you'd like
Models shall be able to save (serialize) and load (deserialize) items as follows :

• Model parameters
• Layers in models

Additional context
Refer to flax serialization logic for more info. It involves serializing and deserializing JAX PyTrees to MessagePack format.

Description
Stax-like layers API refactoring

Solution
Write every layer with apply_fun and init_fun

Feature request description
Regression algorithms missing.

Solution proposed
Implement loss functions like MSE, RMSE, etc.

Additional context
Refer to docs.fast.ai, PyTorch and Keras docs for implementation

Description
Refactor the forward prop code to accommodate different layers in the future.

Solution
Write forward prop code in layers API

Additional context
Refer to https://jax.readthedocs.io/en/latest/index.html

Bug description
APIs too congested in dnet.py script

Additional context
Needs API cleanup from ground up

Description
Implement CNN layers and relevant APIs

Solution
Use LAX APIs

Additional context
Refer to jax.readthedocs.io for more info.