training-mlps

Trains configurable MLPs with Flax NNX for classification and regression tasks.

1|Updated Jan 16, 2026

Search Tags:#jax #model-training #deep-learning #neural-network #flax #mlp

Ready to run

npx skills add https://github.com/yonesuke/skills --skill training-mlps

💬 Prompt for ChatGPT / Claude Chat▼

Please help me install this Agent Skill.
Skill: training-mlps
Source: https://github.com/yonesuke/skills/tree/main/training-mlps
Command: npx skills add https://github.com/yonesuke/skills --skill training-mlps

SYSTEM DOCUMENTATION & REQUIREMENTS

💡 This Skill requires flax, jax, jaxlib, and includes scripts (resource) components.

What problem does it solve?

Provides a reusable, configurable Multi-Layer Perceptron (MLP) backbone for supervised learning, reducing boilerplate in model definition and training.

Core Features & Use Cases

Configurable Architecture: build multiple hidden layers with adjustable size, activation, and output dimension.
Normalization & Regularization: optional layernorm or rmsnorm and dropout support for robust training.
Practical Scenarios: classification and regression tasks requiring a compact feedforward model, quick prototyping, or educational experiments using Flax NNX.

Quick Start

Instantiate MultiLayerPerceptron from the examples, provide input dimensions, and run a forward pass to validate the model.

Frequently Asked Questions about training-mlps

High-intent search queries and answers about installing and using this skill.

FAQPage Schema

How do I train a deep learning MLP model using Flax NNX?▼

To train a deep learning MLP model using Flax NNX, you can instantiate the MultiLayerPerceptron class with configurable depth, width, and activation functions, then execute a forward pass with your supervised learning data.

What is a modular MLP backbone and when do I need it for neural network training?▼

A modular MLP backbone is a reusable Multi-Layer Perceptron architecture for supervised learning. You need it for classification and regression tasks requiring compact feedforward networks, quick prototyping, or educational experiments.

Can I use LayerNorm and dropout for MLP training in JAX?▼

Yes, you can use optional LayerNorm or RMSNorm and dropout handling for robust MLP training in JAX, allowing you to configure normalization and regularization within your feedforward network architecture.

How do I configure hidden layers and output dimensions for a feedforward neural network in JAX?▼

You configure hidden layers and output dimensions for a feedforward neural network in JAX by parameterizing the MultiLayerPerceptron builder, adjusting layer size, activation functions, and the final linear output layer.

What's the best way to build a compact MLP for classification tasks in Flax?▼

The best way to build a compact MLP for classification tasks in Flax is using a modular MultiLayerPerceptron backbone, which reduces boilerplate model definition while providing customizable depth, width, and normalization options.

What are the limitations of using MLPs for deep learning in JAX?▼

MLPs in JAX are limited to small-to-medium network scales and supervised learning tasks like classification and regression, making them unsuitable for large-scale architectures or complex spatial data processing requiring convolutional layers.