DEV Community: Shakudo

A Quick Introduction to JAX

Shakudo — Fri, 21 Jul 2023 15:05:00 +0000

There are enough Python libraries out there that you’ll never understand or use them all. The more pertinent task is choosing the right one for your specific project. At Shakudo it’s pretty common for our team to begin using NumPy or another library, only to figure out halfway through that it’s not effective for our use case.

Shakudo provides data teams with an operating system for data stacks, and we’ve taken a liking to JAX lately for machine learning-and data processing. We’ll explain why in this quick intro to JAX.

What is JAX and what does it do?

Google’s JAX is a high-performance Python package, built to accelerate machine learning research. JAX provides a lightweight API for array-based computing - much like NumPy. It adds a set of composable function transformations, including for automatic differentiation, just-in-time (JIT) compilation, and automated vectorization and parallelization of your code. We’ll talk about those more later on.

JAX is executable on CPU, GPU, or TPU, with minor edits to your code making it easy to speed up big projects in a short amount of time. We’ve seen it used for some really cool projects including protein folding research, robotics control, and physics simulations.

Automatic differentiation is a procedure for computing derivatives that avoids the pitfalls of numerical (expensive and numerically unstable), and symbolic (exponential increase in the number of expressions) differentiation. The automatic differentiation procedure takes a function (program) and simplifies it into a sequence of primitive operations for which the derivative can be easily computed. This procedure is known as backpropagation.

Why use JAX?

Because JAX syntax is so similar to NumPy, with just a few code changes it can be used in projects where NumPy just isn’t cutting it performance-wise, or where you need some extra features that JAX supports. Data-heavy industries including machine learning, blockchain, and other data and compute-heavy use cases benefit from JAX’ improved performance. Maybe you're researching JAX because you’ve hit a wall in terms of scaling your data project.

Beyond speed, JAX is an all around great tool for prototyping because it’s easy to use if you already work with NumPy. It also has powerful features you won’t find in other ML libraries, and a highly familiar syntax for most Python developers.

Tests have shown that JAX can perform up to 8600% faster when used for basic functions - highly valuable for data-heavy application-facing models, or just for getting more machine learning experiments done in a day. Although most real-world applications won’t see this type of speed jump, it does show the potential value of switching.

JAX is capable of these crazy-high speeds for the following reasons:

Vectorization: The method of vectorization enables processing multiple data as a single instruction. This method works for the cases where the same simple operation is applied on the entire data. Since most matrix operations involve applying the same operation on the rows and columns of the matrices, it makes it very amenable to vectorization, providing great speedups for linear algebra computations and machine learning.

JAX allows you to use jax.vmap to automatically generate a vectorized implementation of a function:

import plyvel
def read_chainstate():
    db = plyvel.DB('~/.bitcoin/chainstate', compression=None)

    for key, value in db.iterator(prefix=b'C'):
        tx_key = key[1:]
        print(f"key: {tx_key}"
        print(f"value: {value}")

    db.close()

auto_batch_convolve = jax. vmap(convolve )
auto_batch_convolve(xs, ws)
# DeviceArray([[11., 20., 29. ],
#[11., 20. , 29. ]], dtype=float32)

Code Parallelization: the process of taking a serial code that runs on a single processor and spreading the work across multiple processors. Which means it breaks the problem into smaller pieces so that all data can be processed simultaneously by the computer. This makes the process much more efficient than what it would be by waiting for the solution to one problem to solve the next one.

Automatic differentiation: a set of techniques to evaluate the derivative of a function, by exploiting sequences of elementary arithmetic operations. JAX differentiation is pretty straightforward:

def func(x):
    return x**2
d_func = grad(func)

You can also repeatedly apply grad to get higher order derivatives. That is, we can get the second derivative of func by applying it again on d_func:

d2_func = grad(d_func)

How JAX is built

JAX is built to use Accelerated Linear Algebra (XLA) and Just-in-Time Compilation (JIT). XLA is a domain-specific compiler for linear algebra that fuses together operations, meaning it allows you to skip intermediate results for overall improved speed. JAX uses XLA to compile and run NumPy programs on GPUs and TPUs without changes to your code. It traces your Python code to an intermediate representation, which is then just-in-time compiled.

With JIT, the first time the interpreter runs a method, it gets compiled to machine code so that subsequent executions will run faster. JIT is a simple function:

def funct(x):

    return x * (2 + x)
compiled_funct = jit(funct)

Although it’s a powerful tool, it still doesn't work for every function. You can look to the JIT documentation to understand better about what it can and can’t compile.

How to Install and Import JAX

To install the CPU-only version of JAX, use the following:

pip install --upgrade pip
pip install --upgrade "jax[cuda]"

And that’s it! Now you have the CPU support to test your code. To install GPU support, you’ll need to have CUDA and CuDNN already installed.

Finally, we can import the NumPy interface and the most important JAX functions using:

If you’ve already begun your project using NumPy, you can import it to JAX using the following, and use it to do the same operations as in NumPy:

import jax as jx
import numpy as np # Don't use this
import jax.numpy as jnp # Cool kids do this !

Note that there are two restraints for your NumPy project to work:

You can only use pure functions. If you call your function twice, it has to return the same result, and you can’t do in-place updates of arrays:

array = jnp.zeros((2,2))
array[:, 0] = 1 # Won't work
array = jax.ops.index_update(array, jax.ops.index[:, 0], 1) # Better

Random number generation is explicit and it uses a PRNG key:

key = jx.random.PRNGKey(1)
x = jx.random.normal(key=key, shape=(1000,))

And there you go! You’re off on your first JAX project. If you want to try out all of this in an easily configurable workspace that pre-integrates all the open source tools and data frameworks you want, reach out to our team!

Authors: Stella Wu and Sabrina Aquino

What is Shakudo - The Operating System for Data Stacks

Shakudo — Thu, 20 Jul 2023 21:24:04 +0000

Overview

In the ever-evolving landscape of data science, the founders of Shakudo observed a common roadblock in many organizations - brilliant ideas often never made it into production due to a lack of DevOps engineering resources. Shakudo emerged from the need for a robust solution to this challenge, providing a user-friendly operating system for data stacks.

Why Shakudo?

Shakudo’s mission is to empower data teams by eliminating the complexity of managing their data stack, allowing for the effective implementation of their ideas through:

Effortless Customized Stack

Build a stack that works precisely for your company without worrying about maintenance costs or stability.

Comprehensive Single UI with Automated DevOps

Collaborate on the entire data stack through a single UI, manage access and monitor cost. Reconcile partially overlapping state of tools into a single consistent global state with automated DevOps.

Cloud Cost Optimization

Significantly reduce cloud cost with Shakudo cloud compute management and autoscaling.

Core Platform Components

Shakudo offers a comprehensive data ecosystem that is continually expanding to include a diverse range of best-of-breed tools and technologies. The Platform ensures that your organization is equipped with the most up-to-date data management solutions. Here's how it achieves this goal through its core components:

Sessions: A unified development environment that simplifies the setup process, enabling data teams to start building right away without worrying about environment configuration. This immediate readiness to build enhances productivity and reduces downtime.

Jobs and Services: Whether it's data cleaning in the middle of the night or real-time data updates throughout the day, Jobs ensures your data processing runs smoothly and efficiently. Services take care of deploying features like dashboards, websites, or APIs. Both features offer flexible deployment options, using GIT repositories or pre-built Docker images, and can be scheduled or triggered as required. This orchestration capability means your data teams can manage their workflows more effectively and react quickly to changing requirements.

Shakudo Stack Components: A universe of pre-configured, fully-connected data stacks supporting end-to-end use cases of data and machine learning applications. This feature reduces the time and effort required to setup and connect data stack components, accelerating your path to data insights.

Building a Complete Data Stack Universe

Shakudo is dedicated to building a complete and ever-growing universe of data stack components. It ensures your organization always has access to best-of-breed tools and technologies and the list of integrations is continuously updated to keep up with the fast-paced changes in the data world. You can check out the latest additions on our integrations page.

When to Use Shakudo

Shakudo's platform is adaptable and caters to a diverse range of problems and requirements. Whether you want to quickly develop models, deploy pipelines, monitor model performance, or build data applications, Shakudo provides the right tools and user-friendly environment for your needs.

As your organization grows and your data needs become more complex, Shakudo helps you scale your data infrastructure with ease. For teams eager to explore emerging data technologies or test new tools, it offers a supportive, easy-to-navigate environment.

The Platform designed to support various tools and use cases, including:

Data Engineering: Streamline data transformation development and deployment processes for efficient data management with:

dbt (Data Transformation)
Apache Airflow, Prefect, or Dagster (Pipeline Orchestration)
Airbyte (Data Integration)

Distributed Computing: Manage data larger than memory, optimizing data processing and storage capabilities with:

Dask (Distributed Computing)
Apache Spark (Large-scale Data Processing)
Ray (distributed model training and fine tuning)

Data analytics and Visualization: Enhance data insights and decision-making with advanced analytics and visualization with:

Superset, Cube.dev, Metabase (BI Dashboard)
Streamlit (python based dashboarding tool)

Deployment of Batch Jobs: Automate and manage batch jobs efficiently for improved data processing with:

Apache Airflow, Prefect, or Dagster (Pipeline Orchestration)
Jenkins (Pipeline Automation)

Serving Data Applications and Pipelines: Seamlessly serve and manage data applications and pipelines for better data flow and accessibility with:

Django, FastAPI, or Flask (Web Framework)

Machine Learning Model Training: Train machine learning models effectively, ensuring optimal performance and results with:

TensorFlow, PyTorch, or MXNet (Deep Learning Framework)
Scikit-Learn, XGBoost, or LightGBM (Machine Learning Libraries)
AutoGluon (AutoML)

Machine Learning Model Serving: Deploy and manage machine learning models for production, providing reliable and efficient solutions with:

NVIDIA Triton (Model Serving)
TFServe

Connection to storage and data warehousing: As your organization grows and the volume of data increases, Shakudo can help scale your data infrastructure to accommodate the increasing workload with:

Amazon S3, Google Storage Bucket, or Azure Blob Storage (storage )
Snowflake, BigQuery, or Amazon Redshift (Data Warehouse)

Experimenting with New Data Tools: If your team wants to explore emerging data technologies or test new tools without the burden of DevOps overhead, Shakudo allows for easy experimentation in a flexible environment with:

Langchain
Stanford Alpaca or flan-ul2 (Open source Large Language Models)

In Conclusion

Shakudo is not only an operating system for your data stack but also a strategic partner on your data management journey. The Platform equips your team with the tools needed to innovate and excel in today's data-driven world.

As your organization grows and your data needs evolve, Shakudo scales with you, ensuring your data infrastructure can handle increasing workloads and complexity. To experience the transformative impact of Shakudo's platform, contact our team today.

Author: Sabrina Aquino.
Updated on: July 20, 2023.