import{s as ga,o as Ma,n as Ws}from"../chunks/scheduler.bdbef820.js";import{S as ya,i as Ja,g as m,s as n,r as c,A as fa,h,f as t,c as p,j as da,u as i,x as g,k as ua,y as ba,a as e,v as o,d as j,t as d,w as u}from"../chunks/index.c0aea24a.js";import{T as Ds}from"../chunks/Tip.31005f7d.js";import{C as J}from"../chunks/CodeBlock.6ccca92e.js";import{H as ns,E as Ua}from"../chunks/EditOnGithub.725ee0c1.js";function Ta(f){let l,M=`<code>jax</code> and <code>jaxlib</code> are required to reproduce to code above, so please make sure you
install them as <code>pip install datasets[jax]</code>.`;return{c(){l=m("p"),l.innerHTML=M},l(r){l=h(r,"P",{"data-svelte-h":!0}),g(l)!=="svelte-1i2qrbm"&&(l.innerHTML=M)},m(r,y){e(r,l,y)},p:Ws,d(r){r&&t(l)}}}function wa(f){let l,M='A <a href="/docs/datasets/v3.0.0/en/package_reference/main_classes#datasets.Dataset">Dataset</a> object is a wrapper of an Arrow table, which allows fast reads from arrays in the dataset to JAX arrays.';return{c(){l=m("p"),l.innerHTML=M},l(r){l=h(r,"P",{"data-svelte-h":!0}),g(l)!=="svelte-pnk7ia"&&(l.innerHTML=M)},m(r,y){e(r,l,y)},p:Ws,d(r){r&&t(l)}}}function $a(f){let l,M=`To use the <a href="/docs/datasets/v3.0.0/en/package_reference/main_classes#datasets.Image">Image</a> feature type, you’ll need to install the <code>vision</code> extra as
<code>pip install datasets[vision]</code>.`;return{c(){l=m("p"),l.innerHTML=M},l(r){l=h(r,"P",{"data-svelte-h":!0}),g(l)!=="svelte-1mish4m"&&(l.innerHTML=M)},m(r,y){e(r,l,y)},p:Ws,d(r){r&&t(l)}}}function Ca(f){let l,M=`To use the <a href="/docs/datasets/v3.0.0/en/package_reference/main_classes#datasets.Audio">Audio</a> feature type, you’ll need to install the <code>audio</code> extra as
<code>pip install datasets[audio]</code>.`;return{c(){l=m("p"),l.innerHTML=M},l(r){l=h(r,"P",{"data-svelte-h":!0}),g(l)!=="svelte-1jywyf0"&&(l.innerHTML=M)},m(r,y){e(r,l,y)},p:Ws,d(r){r&&t(l)}}}function Ia(f){let l,M,r,y,$,rs,C,Hs=`This document is a quick introduction to using <code>datasets</code> with JAX, with a particular focus on how to get
<code>jax.Array</code> objects out of our datasets, and how to use them to train JAX models.`,ms,b,hs,I,cs,R,Ls=`By default, datasets return regular Python objects: integers, floats, strings, lists, etc., and
string and binary objects are unchanged, since JAX only supports numbers.`,is,k,Ss="To get JAX arrays (numpy-like) instead, you can set the format of the dataset to <code>jax</code>:",os,Z,js,U,ds,X,Ps=`Note that the exact same procedure applies to <code>DatasetDict</code> objects, so that
when setting the format of a <code>DatasetDict</code> to <code>jax</code>, all the <code>Dataset</code>s there
will be formatted as <code>jax</code>:`,us,v,gs,N,Ks=`Another thing you’ll need to take into consideration is that the formatting is not applied
until you actually access the data. So if you want to get a JAX array out of a dataset,
you’ll need to access the data first, otherwise the format will remain the same.`,Ms,Q,Os=`Finally, to load the data in the device of your choice, you can specify the <code>device</code> argument,
but note that <code>jaxlib.xla_extension.Device</code> is not supported as it’s not serializable with neither
<code>pickle</code> not <code>dill</code>, so you’ll need to use its string identifier instead:`,ys,x,Js,A,sa=`Note that if the <code>device</code> argument is not provided to <code>with_format</code> then it will use the default
device which is <code>jax.devices()[0]</code>.`,fs,q,bs,G,aa="If your dataset consists of N-dimensional arrays, you will see that by default they are considered as the same tensor if the shape is fixed:",Us,_,Ts,E,ws,V,ta=`However this logic often requires slow shape comparisons and data copies.
To avoid this, you must explicitly use the <code>Array</code> feature type and specify the shape of your tensors:`,$s,F,Cs,z,Is,Y,ea='<a href="/docs/datasets/v3.0.0/en/package_reference/main_classes#datasets.ClassLabel">ClassLabel</a> data is properly converted to arrays:',Rs,B,ks,D,la="String and binary objects are unchanged, since JAX only supports numbers.",Zs,W,na='The <a href="/docs/datasets/v3.0.0/en/package_reference/main_classes#datasets.Image">Image</a> and <a href="/docs/datasets/v3.0.0/en/package_reference/main_classes#datasets.Audio">Audio</a> feature types are also supported.',Xs,T,vs,H,Ns,w,Qs,L,xs,S,As,P,pa=`JAX doesn’t have any built-in data loading capabilities, so you’ll need to use a library such
as <a href="https://pytorch.org/" rel="nofollow">PyTorch</a> to load your data using a <code>DataLoader</code> or <a href="https://www.tensorflow.org/" rel="nofollow">TensorFlow</a>
using a <code>tf.data.Dataset</code>. Citing the <a href="https://jax.readthedocs.io/en/latest/notebooks/Neural_Network_and_Data_Loading.html#data-loading-with-pytorch" rel="nofollow">JAX documentation</a> on this topic:
“JAX is laser-focused on program transformations and accelerator-backed NumPy, so we don’t
include data loading or munging in the JAX library. There are already a lot of great data loaders
out there, so let’s just use them instead of reinventing anything. We’ll grab PyTorch’s data loader,
and make a tiny shim to make it work with NumPy arrays.”.`,qs,K,ra=`So that’s the reason why JAX-formatting in <code>datasets</code> is so useful, because it lets you use
any model from the HuggingFace Hub with JAX, without having to worry about the data loading
part.`,Gs,O,_s,ss,ma=`The easiest way to get JAX arrays out of a dataset is to use the <code>with_format(&#39;jax&#39;)</code> method. Lets assume
that we want to train a neural network on the <a href="http://yann.lecun.com/exdb/mnist/" rel="nofollow">MNIST dataset</a> available
at the HuggingFace Hub at <a href="https://huggingface.co/datasets/mnist" rel="nofollow">https://huggingface.co/datasets/mnist</a>.`,Es,as,Vs,ts,ha=`Once the format is set we can feed the dataset to the JAX model in batches using the <code>Dataset.iter()</code>
method:`,Fs,es,zs,ls,Ys,ps,Bs;return $=new ns({props:{title:"Use with JAX",local:"use-with-jax",headingTag:"h1"}}),b=new Ds({props:{$$slots:{default:[Ta]},$$scope:{ctx:f}}}),I=new ns({props:{title:"Dataset format",local:"dataset-format",headingTag:"h2"}}),Z=new J({props:{code:"ZnJvbSUyMGRhdGFzZXRzJTIwaW1wb3J0JTIwRGF0YXNldCUwQWRhdGElMjAlM0QlMjAlNUIlNUIxJTJDJTIwMiU1RCUyQyUyMCU1QjMlMkMlMjA0JTVEJTVEJTBBZHMlMjAlM0QlMjBEYXRhc2V0LmZyb21fZGljdCglN0IlMjJkYXRhJTIyJTNBJTIwZGF0YSU3RCklMEFkcyUyMCUzRCUyMGRzLndpdGhfZm9ybWF0KCUyMmpheCUyMiklMEFkcyU1QjAlNUQlMEFkcyU1QiUzQTIlNUQ=",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> Dataset
<span class="hljs-meta">&gt;&gt;&gt; </span>data = [[<span class="hljs-number">1</span>, <span class="hljs-number">2</span>], [<span class="hljs-number">3</span>, <span class="hljs-number">4</span>]]
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = Dataset.from_dict({<span class="hljs-string">&quot;data&quot;</span>: data})
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = ds.with_format(<span class="hljs-string">&quot;jax&quot;</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[<span class="hljs-number">0</span>]
{<span class="hljs-string">&#x27;data&#x27;</span>: DeviceArray([<span class="hljs-number">1</span>, <span class="hljs-number">2</span>], dtype=int32)}
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[:<span class="hljs-number">2</span>]
{<span class="hljs-string">&#x27;data&#x27;</span>: DeviceArray([
    [<span class="hljs-number">1</span>, <span class="hljs-number">2</span>],
    [<span class="hljs-number">3</span>, <span class="hljs-number">4</span>]], dtype=int32)}`,wrap:!1}}),U=new Ds({props:{$$slots:{default:[wa]},$$scope:{ctx:f}}}),v=new J({props:{code:"ZnJvbSUyMGRhdGFzZXRzJTIwaW1wb3J0JTIwRGF0YXNldERpY3QlMEFkYXRhJTIwJTNEJTIwJTdCJTIydHJhaW4lMjIlM0ElMjAlN0IlMjJkYXRhJTIyJTNBJTIwJTVCJTVCMSUyQyUyMDIlNUQlMkMlMjAlNUIzJTJDJTIwNCU1RCU1RCU3RCUyQyUyMCUyMnRlc3QlMjIlM0ElMjAlN0IlMjJkYXRhJTIyJTNBJTIwJTVCJTVCNSUyQyUyMDYlNUQlMkMlMjAlNUI3JTJDJTIwOCU1RCU1RCU3RCU3RCUwQWRkcyUyMCUzRCUyMERhdGFzZXREaWN0LmZyb21fZGljdChkYXRhKSUwQWRkcyUyMCUzRCUyMGRkcy53aXRoX2Zvcm1hdCglMjJqYXglMjIpJTBBZGRzJTVCJTIydHJhaW4lMjIlNUQlNUIlM0EyJTVE",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> DatasetDict
<span class="hljs-meta">&gt;&gt;&gt; </span>data = {<span class="hljs-string">&quot;train&quot;</span>: {<span class="hljs-string">&quot;data&quot;</span>: [[<span class="hljs-number">1</span>, <span class="hljs-number">2</span>], [<span class="hljs-number">3</span>, <span class="hljs-number">4</span>]]}, <span class="hljs-string">&quot;test&quot;</span>: {<span class="hljs-string">&quot;data&quot;</span>: [[<span class="hljs-number">5</span>, <span class="hljs-number">6</span>], [<span class="hljs-number">7</span>, <span class="hljs-number">8</span>]]}}
<span class="hljs-meta">&gt;&gt;&gt; </span>dds = DatasetDict.from_dict(data)
<span class="hljs-meta">&gt;&gt;&gt; </span>dds = dds.with_format(<span class="hljs-string">&quot;jax&quot;</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span>dds[<span class="hljs-string">&quot;train&quot;</span>][:<span class="hljs-number">2</span>]
{<span class="hljs-string">&#x27;data&#x27;</span>: DeviceArray([
    [<span class="hljs-number">1</span>, <span class="hljs-number">2</span>],
    [<span class="hljs-number">3</span>, <span class="hljs-number">4</span>]], dtype=int32)}`,wrap:!1}}),x=new J({props:{code:"aW1wb3J0JTIwamF4JTBBZnJvbSUyMGRhdGFzZXRzJTIwaW1wb3J0JTIwRGF0YXNldCUwQWRhdGElMjAlM0QlMjAlNUIlNUIxJTJDJTIwMiU1RCUyQyUyMCU1QjMlMkMlMjA0JTVEJTVEJTBBZHMlMjAlM0QlMjBEYXRhc2V0LmZyb21fZGljdCglN0IlMjJkYXRhJTIyJTNBJTIwZGF0YSU3RCklMEFkZXZpY2UlMjAlM0QlMjBzdHIoamF4LmRldmljZXMoKSU1QjAlNUQpJTIwJTIwJTIzJTIwTm90JTIwY2FzdGluZyUyMHRvJTIwJTYwc3RyJTYwJTIwYmVmb3JlJTIwcGFzc2luZyUyMGl0JTIwdG8lMjAlNjB3aXRoX2Zvcm1hdCU2MCUyMHdpbGwlMjByYWlzZSUyMGElMjAlNjBWYWx1ZUVycm9yJTYwJTBBZHMlMjAlM0QlMjBkcy53aXRoX2Zvcm1hdCglMjJqYXglMjIlMkMlMjBkZXZpY2UlM0RkZXZpY2UpJTBBZHMlNUIwJTVEJTBBZHMlNUIwJTVEJTVCJTIyZGF0YSUyMiU1RC5kZXZpY2UoKSUwQWFzc2VydCUyMGRzJTVCMCU1RCU1QiUyMmRhdGElMjIlNUQuZGV2aWNlKCklMjAlM0QlM0QlMjBqYXguZGV2aWNlcygpJTVCMCU1RA==",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">import</span> jax
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> Dataset
<span class="hljs-meta">&gt;&gt;&gt; </span>data = [[<span class="hljs-number">1</span>, <span class="hljs-number">2</span>], [<span class="hljs-number">3</span>, <span class="hljs-number">4</span>]]
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = Dataset.from_dict({<span class="hljs-string">&quot;data&quot;</span>: data})
<span class="hljs-meta">&gt;&gt;&gt; </span>device = <span class="hljs-built_in">str</span>(jax.devices()[<span class="hljs-number">0</span>])  <span class="hljs-comment"># Not casting to \`str\` before passing it to \`with_format\` will raise a \`ValueError\`</span>
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = ds.with_format(<span class="hljs-string">&quot;jax&quot;</span>, device=device)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[<span class="hljs-number">0</span>]
{<span class="hljs-string">&#x27;data&#x27;</span>: DeviceArray([<span class="hljs-number">1</span>, <span class="hljs-number">2</span>], dtype=int32)}
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[<span class="hljs-number">0</span>][<span class="hljs-string">&quot;data&quot;</span>].device()
TFRT_CPU_0
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">assert</span> ds[<span class="hljs-number">0</span>][<span class="hljs-string">&quot;data&quot;</span>].device() == jax.devices()[<span class="hljs-number">0</span>]
<span class="hljs-literal">True</span>`,wrap:!1}}),q=new ns({props:{title:"N-dimensional arrays",local:"n-dimensional-arrays",headingTag:"h3"}}),_=new J({props:{code:"ZnJvbSUyMGRhdGFzZXRzJTIwaW1wb3J0JTIwRGF0YXNldCUwQWRhdGElMjAlM0QlMjAlNUIlNUIlNUIxJTJDJTIwMiU1RCUyQyU1QjMlMkMlMjA0JTVEJTVEJTJDJTIwJTVCJTVCNSUyQyUyMDYlNUQlMkMlNUI3JTJDJTIwOCU1RCU1RCU1RCUyMCUyMCUyMyUyMGZpeGVkJTIwc2hhcGUlMEFkcyUyMCUzRCUyMERhdGFzZXQuZnJvbV9kaWN0KCU3QiUyMmRhdGElMjIlM0ElMjBkYXRhJTdEKSUwQWRzJTIwJTNEJTIwZHMud2l0aF9mb3JtYXQoJTIyamF4JTIyKSUwQWRzJTVCMCU1RA==",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> Dataset
<span class="hljs-meta">&gt;&gt;&gt; </span>data = [[[<span class="hljs-number">1</span>, <span class="hljs-number">2</span>],[<span class="hljs-number">3</span>, <span class="hljs-number">4</span>]], [[<span class="hljs-number">5</span>, <span class="hljs-number">6</span>],[<span class="hljs-number">7</span>, <span class="hljs-number">8</span>]]]  <span class="hljs-comment"># fixed shape</span>
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = Dataset.from_dict({<span class="hljs-string">&quot;data&quot;</span>: data})
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = ds.with_format(<span class="hljs-string">&quot;jax&quot;</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[<span class="hljs-number">0</span>]
{<span class="hljs-string">&#x27;data&#x27;</span>: Array([[<span class="hljs-number">1</span>, <span class="hljs-number">2</span>],
        [<span class="hljs-number">3</span>, <span class="hljs-number">4</span>]], dtype=int32)}`,wrap:!1}}),E=new J({props:{code:"ZnJvbSUyMGRhdGFzZXRzJTIwaW1wb3J0JTIwRGF0YXNldCUwQWRhdGElMjAlM0QlMjAlNUIlNUIlNUIxJTJDJTIwMiU1RCUyQyU1QjMlNUQlNUQlMkMlMjAlNUIlNUI0JTJDJTIwNSUyQyUyMDYlNUQlMkMlNUI3JTJDJTIwOCU1RCU1RCU1RCUyMCUyMCUyMyUyMHZhcnlpbmclMjBzaGFwZSUwQWRzJTIwJTNEJTIwRGF0YXNldC5mcm9tX2RpY3QoJTdCJTIyZGF0YSUyMiUzQSUyMGRhdGElN0QpJTBBZHMlMjAlM0QlMjBkcy53aXRoX2Zvcm1hdCglMjJqYXglMjIpJTBBZHMlNUIwJTVE",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> Dataset
<span class="hljs-meta">&gt;&gt;&gt; </span>data = [[[<span class="hljs-number">1</span>, <span class="hljs-number">2</span>],[<span class="hljs-number">3</span>]], [[<span class="hljs-number">4</span>, <span class="hljs-number">5</span>, <span class="hljs-number">6</span>],[<span class="hljs-number">7</span>, <span class="hljs-number">8</span>]]]  <span class="hljs-comment"># varying shape</span>
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = Dataset.from_dict({<span class="hljs-string">&quot;data&quot;</span>: data})
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = ds.with_format(<span class="hljs-string">&quot;jax&quot;</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[<span class="hljs-number">0</span>]
{<span class="hljs-string">&#x27;data&#x27;</span>: [Array([<span class="hljs-number">1</span>, <span class="hljs-number">2</span>], dtype=int32), Array([<span class="hljs-number">3</span>], dtype=int32)]}`,wrap:!1}}),F=new J({props:{code:"ZnJvbSUyMGRhdGFzZXRzJTIwaW1wb3J0JTIwRGF0YXNldCUyQyUyMEZlYXR1cmVzJTJDJTIwQXJyYXkyRCUwQWRhdGElMjAlM0QlMjAlNUIlNUIlNUIxJTJDJTIwMiU1RCUyQyU1QjMlMkMlMjA0JTVEJTVEJTJDJTVCJTVCNSUyQyUyMDYlNUQlMkMlNUI3JTJDJTIwOCU1RCU1RCU1RCUwQWZlYXR1cmVzJTIwJTNEJTIwRmVhdHVyZXMoJTdCJTIyZGF0YSUyMiUzQSUyMEFycmF5MkQoc2hhcGUlM0QoMiUyQyUyMDIpJTJDJTIwZHR5cGUlM0QnaW50MzInKSU3RCklMEFkcyUyMCUzRCUyMERhdGFzZXQuZnJvbV9kaWN0KCU3QiUyMmRhdGElMjIlM0ElMjBkYXRhJTdEJTJDJTIwZmVhdHVyZXMlM0RmZWF0dXJlcyklMEFkcyUyMCUzRCUyMGRzLndpdGhfZm9ybWF0KCUyMnRvcmNoJTIyKSUwQWRzJTVCMCU1RCUwQWRzJTVCJTNBMiU1RA==",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> Dataset, Features, Array2D
<span class="hljs-meta">&gt;&gt;&gt; </span>data = [[[<span class="hljs-number">1</span>, <span class="hljs-number">2</span>],[<span class="hljs-number">3</span>, <span class="hljs-number">4</span>]],[[<span class="hljs-number">5</span>, <span class="hljs-number">6</span>],[<span class="hljs-number">7</span>, <span class="hljs-number">8</span>]]]
<span class="hljs-meta">&gt;&gt;&gt; </span>features = Features({<span class="hljs-string">&quot;data&quot;</span>: Array2D(shape=(<span class="hljs-number">2</span>, <span class="hljs-number">2</span>), dtype=<span class="hljs-string">&#x27;int32&#x27;</span>)})
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = Dataset.from_dict({<span class="hljs-string">&quot;data&quot;</span>: data}, features=features)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = ds.with_format(<span class="hljs-string">&quot;torch&quot;</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[<span class="hljs-number">0</span>]
{<span class="hljs-string">&#x27;data&#x27;</span>: Array([[<span class="hljs-number">1</span>, <span class="hljs-number">2</span>],
        [<span class="hljs-number">3</span>, <span class="hljs-number">4</span>]], dtype=int32)}
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[:<span class="hljs-number">2</span>]
{<span class="hljs-string">&#x27;data&#x27;</span>: Array([[[<span class="hljs-number">1</span>, <span class="hljs-number">2</span>],
         [<span class="hljs-number">3</span>, <span class="hljs-number">4</span>]],
 
        [[<span class="hljs-number">5</span>, <span class="hljs-number">6</span>],
         [<span class="hljs-number">7</span>, <span class="hljs-number">8</span>]]], dtype=int32)}`,wrap:!1}}),z=new ns({props:{title:"Other feature types",local:"other-feature-types",headingTag:"h3"}}),B=new J({props:{code:"ZnJvbSUyMGRhdGFzZXRzJTIwaW1wb3J0JTIwRGF0YXNldCUyQyUyMEZlYXR1cmVzJTJDJTIwQ2xhc3NMYWJlbCUwQWxhYmVscyUyMCUzRCUyMCU1QjAlMkMlMjAwJTJDJTIwMSU1RCUwQWZlYXR1cmVzJTIwJTNEJTIwRmVhdHVyZXMoJTdCJTIybGFiZWwlMjIlM0ElMjBDbGFzc0xhYmVsKG5hbWVzJTNEJTVCJTIybmVnYXRpdmUlMjIlMkMlMjAlMjJwb3NpdGl2ZSUyMiU1RCklN0QpJTBBZHMlMjAlM0QlMjBEYXRhc2V0LmZyb21fZGljdCglN0IlMjJsYWJlbCUyMiUzQSUyMGxhYmVscyU3RCUyQyUyMGZlYXR1cmVzJTNEZmVhdHVyZXMpJTBBZHMlMjAlM0QlMjBkcy53aXRoX2Zvcm1hdCglMjJqYXglMjIpJTBBZHMlNUIlM0EzJTVE",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> Dataset, Features, ClassLabel
<span class="hljs-meta">&gt;&gt;&gt; </span>labels = [<span class="hljs-number">0</span>, <span class="hljs-number">0</span>, <span class="hljs-number">1</span>]
<span class="hljs-meta">&gt;&gt;&gt; </span>features = Features({<span class="hljs-string">&quot;label&quot;</span>: ClassLabel(names=[<span class="hljs-string">&quot;negative&quot;</span>, <span class="hljs-string">&quot;positive&quot;</span>])})
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = Dataset.from_dict({<span class="hljs-string">&quot;label&quot;</span>: labels}, features=features)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = ds.with_format(<span class="hljs-string">&quot;jax&quot;</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[:<span class="hljs-number">3</span>]
{<span class="hljs-string">&#x27;label&#x27;</span>: DeviceArray([<span class="hljs-number">0</span>, <span class="hljs-number">0</span>, <span class="hljs-number">1</span>], dtype=int32)}`,wrap:!1}}),T=new Ds({props:{$$slots:{default:[$a]},$$scope:{ctx:f}}}),H=new J({props:{code:"ZnJvbSUyMGRhdGFzZXRzJTIwaW1wb3J0JTIwRGF0YXNldCUyQyUyMEZlYXR1cmVzJTJDJTIwSW1hZ2UlMEFpbWFnZXMlMjAlM0QlMjAlNUIlMjJwYXRoJTJGdG8lMkZpbWFnZS5wbmclMjIlNUQlMjAqJTIwMTAlMEFmZWF0dXJlcyUyMCUzRCUyMEZlYXR1cmVzKCU3QiUyMmltYWdlJTIyJTNBJTIwSW1hZ2UoKSU3RCklMEFkcyUyMCUzRCUyMERhdGFzZXQuZnJvbV9kaWN0KCU3QiUyMmltYWdlJTIyJTNBJTIwaW1hZ2VzJTdEJTJDJTIwZmVhdHVyZXMlM0RmZWF0dXJlcyklMEFkcyUyMCUzRCUyMGRzLndpdGhfZm9ybWF0KCUyMmpheCUyMiklMEFkcyU1QjAlNUQlNUIlMjJpbWFnZSUyMiU1RC5zaGFwZSUwQWRzJTVCMCU1RCUwQWRzJTVCJTNBMiU1RCU1QiUyMmltYWdlJTIyJTVELnNoYXBlJTBBZHMlNUIlM0EyJTVE",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> Dataset, Features, Image
<span class="hljs-meta">&gt;&gt;&gt; </span>images = [<span class="hljs-string">&quot;path/to/image.png&quot;</span>] * <span class="hljs-number">10</span>
<span class="hljs-meta">&gt;&gt;&gt; </span>features = Features({<span class="hljs-string">&quot;image&quot;</span>: Image()})
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = Dataset.from_dict({<span class="hljs-string">&quot;image&quot;</span>: images}, features=features)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = ds.with_format(<span class="hljs-string">&quot;jax&quot;</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[<span class="hljs-number">0</span>][<span class="hljs-string">&quot;image&quot;</span>].shape
(<span class="hljs-number">512</span>, <span class="hljs-number">512</span>, <span class="hljs-number">3</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[<span class="hljs-number">0</span>]
{<span class="hljs-string">&#x27;image&#x27;</span>: DeviceArray([[[ <span class="hljs-number">255</span>, <span class="hljs-number">255</span>, <span class="hljs-number">255</span>],
              [ <span class="hljs-number">255</span>, <span class="hljs-number">255</span>, <span class="hljs-number">255</span>],
              ...,
              [ <span class="hljs-number">255</span>, <span class="hljs-number">255</span>, <span class="hljs-number">255</span>],
              [ <span class="hljs-number">255</span>, <span class="hljs-number">255</span>, <span class="hljs-number">255</span>]]], dtype=uint8)}
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[:<span class="hljs-number">2</span>][<span class="hljs-string">&quot;image&quot;</span>].shape
(<span class="hljs-number">2</span>, <span class="hljs-number">512</span>, <span class="hljs-number">512</span>, <span class="hljs-number">3</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[:<span class="hljs-number">2</span>]
{<span class="hljs-string">&#x27;image&#x27;</span>: DeviceArray([[[[ <span class="hljs-number">255</span>, <span class="hljs-number">255</span>, <span class="hljs-number">255</span>],
              [ <span class="hljs-number">255</span>, <span class="hljs-number">255</span>, <span class="hljs-number">255</span>],
              ...,
              [ <span class="hljs-number">255</span>, <span class="hljs-number">255</span>, <span class="hljs-number">255</span>],
              [ <span class="hljs-number">255</span>, <span class="hljs-number">255</span>, <span class="hljs-number">255</span>]]]], dtype=uint8)}`,wrap:!1}}),w=new Ds({props:{$$slots:{default:[Ca]},$$scope:{ctx:f}}}),L=new J({props:{code:"ZnJvbSUyMGRhdGFzZXRzJTIwaW1wb3J0JTIwRGF0YXNldCUyQyUyMEZlYXR1cmVzJTJDJTIwQXVkaW8lMEFhdWRpbyUyMCUzRCUyMCU1QiUyMnBhdGglMkZ0byUyRmF1ZGlvLndhdiUyMiU1RCUyMColMjAxMCUwQWZlYXR1cmVzJTIwJTNEJTIwRmVhdHVyZXMoJTdCJTIyYXVkaW8lMjIlM0ElMjBBdWRpbygpJTdEKSUwQWRzJTIwJTNEJTIwRGF0YXNldC5mcm9tX2RpY3QoJTdCJTIyYXVkaW8lMjIlM0ElMjBhdWRpbyU3RCUyQyUyMGZlYXR1cmVzJTNEZmVhdHVyZXMpJTBBZHMlMjAlM0QlMjBkcy53aXRoX2Zvcm1hdCglMjJqYXglMjIpJTBBZHMlNUIwJTVEJTVCJTIyYXVkaW8lMjIlNUQlNUIlMjJhcnJheSUyMiU1RCUwQWRzJTVCMCU1RCU1QiUyMmF1ZGlvJTIyJTVEJTVCJTIyc2FtcGxpbmdfcmF0ZSUyMiU1RA==",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> Dataset, Features, Audio
<span class="hljs-meta">&gt;&gt;&gt; </span>audio = [<span class="hljs-string">&quot;path/to/audio.wav&quot;</span>] * <span class="hljs-number">10</span>
<span class="hljs-meta">&gt;&gt;&gt; </span>features = Features({<span class="hljs-string">&quot;audio&quot;</span>: Audio()})
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = Dataset.from_dict({<span class="hljs-string">&quot;audio&quot;</span>: audio}, features=features)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = ds.with_format(<span class="hljs-string">&quot;jax&quot;</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[<span class="hljs-number">0</span>][<span class="hljs-string">&quot;audio&quot;</span>][<span class="hljs-string">&quot;array&quot;</span>]
DeviceArray([-<span class="hljs-number">0.059021</span>  , -<span class="hljs-number">0.03894043</span>, -<span class="hljs-number">0.00735474</span>, ...,  <span class="hljs-number">0.0133667</span> ,
              <span class="hljs-number">0.01809692</span>,  <span class="hljs-number">0.00268555</span>], dtype=float32)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[<span class="hljs-number">0</span>][<span class="hljs-string">&quot;audio&quot;</span>][<span class="hljs-string">&quot;sampling_rate&quot;</span>]
DeviceArray(<span class="hljs-number">44100</span>, dtype=int32, weak_type=<span class="hljs-literal">True</span>)`,wrap:!1}}),S=new ns({props:{title:"Data loading",local:"data-loading",headingTag:"h2"}}),O=new ns({props:{title:"Using with_format('jax')",local:"using-withformatjax",headingTag:"h3"}}),as=new J({props:{code:"ZnJvbSUyMGRhdGFzZXRzJTIwaW1wb3J0JTIwbG9hZF9kYXRhc2V0JTBBZHMlMjAlM0QlMjBsb2FkX2RhdGFzZXQoJTIybW5pc3QlMjIpJTBBZHMlMjAlM0QlMjBkcy53aXRoX2Zvcm1hdCglMjJqYXglMjIpJTBBZHMlNUIlMjJ0cmFpbiUyMiU1RCU1QjAlNUQ=",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> load_dataset
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = load_dataset(<span class="hljs-string">&quot;mnist&quot;</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds = ds.with_format(<span class="hljs-string">&quot;jax&quot;</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span>ds[<span class="hljs-string">&quot;train&quot;</span>][<span class="hljs-number">0</span>]
{<span class="hljs-string">&#x27;image&#x27;</span>: DeviceArray([[  <span class="hljs-number">0</span>,   <span class="hljs-number">0</span>,   <span class="hljs-number">0</span>, ...],
                       [  <span class="hljs-number">0</span>,   <span class="hljs-number">0</span>,   <span class="hljs-number">0</span>, ...],
                       ...,
                       [  <span class="hljs-number">0</span>,   <span class="hljs-number">0</span>,   <span class="hljs-number">0</span>, ...],
                       [  <span class="hljs-number">0</span>,   <span class="hljs-number">0</span>,   <span class="hljs-number">0</span>, ...]], dtype=uint8),
 <span class="hljs-string">&#x27;label&#x27;</span>: DeviceArray(<span class="hljs-number">5</span>, dtype=int32)}`,wrap:!1}}),es=new J({props:{code:"Zm9yJTIwZXBvY2glMjBpbiUyMHJhbmdlKGVwb2NocyklM0ElMEElMjAlMjAlMjAlMjBmb3IlMjBiYXRjaCUyMGluJTIwZHMlNUIlMjJ0cmFpbiUyMiU1RC5pdGVyKGJhdGNoX3NpemUlM0QzMiklM0ElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjB4JTJDJTIweSUyMCUzRCUyMGJhdGNoJTVCJTIyaW1hZ2UlMjIlNUQlMkMlMjBiYXRjaCU1QiUyMmxhYmVsJTIyJTVEJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwLi4u",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">for</span> epoch <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(epochs):
<span class="hljs-meta">... </span>    <span class="hljs-keyword">for</span> batch <span class="hljs-keyword">in</span> ds[<span class="hljs-string">&quot;train&quot;</span>].<span class="hljs-built_in">iter</span>(batch_size=<span class="hljs-number">32</span>):
<span class="hljs-meta">... </span>        x, y = batch[<span class="hljs-string">&quot;image&quot;</span>], batch[<span class="hljs-string">&quot;label&quot;</span>]
<span class="hljs-meta">... </span>        ...`,wrap:!1}}),ls=new Ua({props:{source:"https://github.com/huggingface/datasets/blob/main/docs/source/use_with_jax.mdx"}}),{c(){l=m("meta"),M=n(),r=m("p"),y=n(),c($.$$.fragment),rs=n(),C=m("p"),C.innerHTML=Hs,ms=n(),c(b.$$.fragment),hs=n(),c(I.$$.fragment),cs=n(),R=m("p"),R.textContent=Ls,is=n(),k=m("p"),k.innerHTML=Ss,os=n(),c(Z.$$.fragment),js=n(),c(U.$$.fragment),ds=n(),X=m("p"),X.innerHTML=Ps,us=n(),c(v.$$.fragment),gs=n(),N=m("p"),N.textContent=Ks,Ms=n(),Q=m("p"),Q.innerHTML=Os,ys=n(),c(x.$$.fragment),Js=n(),A=m("p"),A.innerHTML=sa,fs=n(),c(q.$$.fragment),bs=n(),G=m("p"),G.textContent=aa,Us=n(),c(_.$$.fragment),Ts=n(),c(E.$$.fragment),ws=n(),V=m("p"),V.innerHTML=ta,$s=n(),c(F.$$.fragment),Cs=n(),c(z.$$.fragment),Is=n(),Y=m("p"),Y.innerHTML=ea,Rs=n(),c(B.$$.fragment),ks=n(),D=m("p"),D.textContent=la,Zs=n(),W=m("p"),W.innerHTML=na,Xs=n(),c(T.$$.fragment),vs=n(),c(H.$$.fragment),Ns=n(),c(w.$$.fragment),Qs=n(),c(L.$$.fragment),xs=n(),c(S.$$.fragment),As=n(),P=m("p"),P.innerHTML=pa,qs=n(),K=m("p"),K.innerHTML=ra,Gs=n(),c(O.$$.fragment),_s=n(),ss=m("p"),ss.innerHTML=ma,Es=n(),c(as.$$.fragment),Vs=n(),ts=m("p"),ts.innerHTML=ha,Fs=n(),c(es.$$.fragment),zs=n(),c(ls.$$.fragment),Ys=n(),ps=m("p"),this.h()},l(s){const a=fa("svelte-u9bgzb",document.head);l=h(a,"META",{name:!0,content:!0}),a.forEach(t),M=p(s),r=h(s,"P",{}),da(r).forEach(t),y=p(s),i($.$$.fragment,s),rs=p(s),C=h(s,"P",{"data-svelte-h":!0}),g(C)!=="svelte-1a2kcgx"&&(C.innerHTML=Hs),ms=p(s),i(b.$$.fragment,s),hs=p(s),i(I.$$.fragment,s),cs=p(s),R=h(s,"P",{"data-svelte-h":!0}),g(R)!=="svelte-1ix49d8"&&(R.textContent=Ls),is=p(s),k=h(s,"P",{"data-svelte-h":!0}),g(k)!=="svelte-1upv66g"&&(k.innerHTML=Ss),os=p(s),i(Z.$$.fragment,s),js=p(s),i(U.$$.fragment,s),ds=p(s),X=h(s,"P",{"data-svelte-h":!0}),g(X)!=="svelte-1b7i338"&&(X.innerHTML=Ps),us=p(s),i(v.$$.fragment,s),gs=p(s),N=h(s,"P",{"data-svelte-h":!0}),g(N)!=="svelte-1c29vp0"&&(N.textContent=Ks),Ms=p(s),Q=h(s,"P",{"data-svelte-h":!0}),g(Q)!=="svelte-1ooxy7c"&&(Q.innerHTML=Os),ys=p(s),i(x.$$.fragment,s),Js=p(s),A=h(s,"P",{"data-svelte-h":!0}),g(A)!=="svelte-hj2cy7"&&(A.innerHTML=sa),fs=p(s),i(q.$$.fragment,s),bs=p(s),G=h(s,"P",{"data-svelte-h":!0}),g(G)!=="svelte-smjp9l"&&(G.textContent=aa),Us=p(s),i(_.$$.fragment,s),Ts=p(s),i(E.$$.fragment,s),ws=p(s),V=h(s,"P",{"data-svelte-h":!0}),g(V)!=="svelte-1gw41y9"&&(V.innerHTML=ta),$s=p(s),i(F.$$.fragment,s),Cs=p(s),i(z.$$.fragment,s),Is=p(s),Y=h(s,"P",{"data-svelte-h":!0}),g(Y)!=="svelte-uo9ve5"&&(Y.innerHTML=ea),Rs=p(s),i(B.$$.fragment,s),ks=p(s),D=h(s,"P",{"data-svelte-h":!0}),g(D)!=="svelte-gkri0z"&&(D.textContent=la),Zs=p(s),W=h(s,"P",{"data-svelte-h":!0}),g(W)!=="svelte-h9k3oa"&&(W.innerHTML=na),Xs=p(s),i(T.$$.fragment,s),vs=p(s),i(H.$$.fragment,s),Ns=p(s),i(w.$$.fragment,s),Qs=p(s),i(L.$$.fragment,s),xs=p(s),i(S.$$.fragment,s),As=p(s),P=h(s,"P",{"data-svelte-h":!0}),g(P)!=="svelte-1msw6w0"&&(P.innerHTML=pa),qs=p(s),K=h(s,"P",{"data-svelte-h":!0}),g(K)!=="svelte-1xmbhz7"&&(K.innerHTML=ra),Gs=p(s),i(O.$$.fragment,s),_s=p(s),ss=h(s,"P",{"data-svelte-h":!0}),g(ss)!=="svelte-1pw5xoa"&&(ss.innerHTML=ma),Es=p(s),i(as.$$.fragment,s),Vs=p(s),ts=h(s,"P",{"data-svelte-h":!0}),g(ts)!=="svelte-lnmbh3"&&(ts.innerHTML=ha),Fs=p(s),i(es.$$.fragment,s),zs=p(s),i(ls.$$.fragment,s),Ys=p(s),ps=h(s,"P",{}),da(ps).forEach(t),this.h()},h(){ua(l,"name","hf:doc:metadata"),ua(l,"content",Ra)},m(s,a){ba(document.head,l),e(s,M,a),e(s,r,a),e(s,y,a),o($,s,a),e(s,rs,a),e(s,C,a),e(s,ms,a),o(b,s,a),e(s,hs,a),o(I,s,a),e(s,cs,a),e(s,R,a),e(s,is,a),e(s,k,a),e(s,os,a),o(Z,s,a),e(s,js,a),o(U,s,a),e(s,ds,a),e(s,X,a),e(s,us,a),o(v,s,a),e(s,gs,a),e(s,N,a),e(s,Ms,a),e(s,Q,a),e(s,ys,a),o(x,s,a),e(s,Js,a),e(s,A,a),e(s,fs,a),o(q,s,a),e(s,bs,a),e(s,G,a),e(s,Us,a),o(_,s,a),e(s,Ts,a),o(E,s,a),e(s,ws,a),e(s,V,a),e(s,$s,a),o(F,s,a),e(s,Cs,a),o(z,s,a),e(s,Is,a),e(s,Y,a),e(s,Rs,a),o(B,s,a),e(s,ks,a),e(s,D,a),e(s,Zs,a),e(s,W,a),e(s,Xs,a),o(T,s,a),e(s,vs,a),o(H,s,a),e(s,Ns,a),o(w,s,a),e(s,Qs,a),o(L,s,a),e(s,xs,a),o(S,s,a),e(s,As,a),e(s,P,a),e(s,qs,a),e(s,K,a),e(s,Gs,a),o(O,s,a),e(s,_s,a),e(s,ss,a),e(s,Es,a),o(as,s,a),e(s,Vs,a),e(s,ts,a),e(s,Fs,a),o(es,s,a),e(s,zs,a),o(ls,s,a),e(s,Ys,a),e(s,ps,a),Bs=!0},p(s,[a]){const ca={};a&2&&(ca.$$scope={dirty:a,ctx:s}),b.$set(ca);const ia={};a&2&&(ia.$$scope={dirty:a,ctx:s}),U.$set(ia);const oa={};a&2&&(oa.$$scope={dirty:a,ctx:s}),T.$set(oa);const ja={};a&2&&(ja.$$scope={dirty:a,ctx:s}),w.$set(ja)},i(s){Bs||(j($.$$.fragment,s),j(b.$$.fragment,s),j(I.$$.fragment,s),j(Z.$$.fragment,s),j(U.$$.fragment,s),j(v.$$.fragment,s),j(x.$$.fragment,s),j(q.$$.fragment,s),j(_.$$.fragment,s),j(E.$$.fragment,s),j(F.$$.fragment,s),j(z.$$.fragment,s),j(B.$$.fragment,s),j(T.$$.fragment,s),j(H.$$.fragment,s),j(w.$$.fragment,s),j(L.$$.fragment,s),j(S.$$.fragment,s),j(O.$$.fragment,s),j(as.$$.fragment,s),j(es.$$.fragment,s),j(ls.$$.fragment,s),Bs=!0)},o(s){d($.$$.fragment,s),d(b.$$.fragment,s),d(I.$$.fragment,s),d(Z.$$.fragment,s),d(U.$$.fragment,s),d(v.$$.fragment,s),d(x.$$.fragment,s),d(q.$$.fragment,s),d(_.$$.fragment,s),d(E.$$.fragment,s),d(F.$$.fragment,s),d(z.$$.fragment,s),d(B.$$.fragment,s),d(T.$$.fragment,s),d(H.$$.fragment,s),d(w.$$.fragment,s),d(L.$$.fragment,s),d(S.$$.fragment,s),d(O.$$.fragment,s),d(as.$$.fragment,s),d(es.$$.fragment,s),d(ls.$$.fragment,s),Bs=!1},d(s){s&&(t(M),t(r),t(y),t(rs),t(C),t(ms),t(hs),t(cs),t(R),t(is),t(k),t(os),t(js),t(ds),t(X),t(us),t(gs),t(N),t(Ms),t(Q),t(ys),t(Js),t(A),t(fs),t(bs),t(G),t(Us),t(Ts),t(ws),t(V),t($s),t(Cs),t(Is),t(Y),t(Rs),t(ks),t(D),t(Zs),t(W),t(Xs),t(vs),t(Ns),t(Qs),t(xs),t(As),t(P),t(qs),t(K),t(Gs),t(_s),t(ss),t(Es),t(Vs),t(ts),t(Fs),t(zs),t(Ys),t(ps)),t(l),u($,s),u(b,s),u(I,s),u(Z,s),u(U,s),u(v,s),u(x,s),u(q,s),u(_,s),u(E,s),u(F,s),u(z,s),u(B,s),u(T,s),u(H,s),u(w,s),u(L,s),u(S,s),u(O,s),u(as,s),u(es,s),u(ls,s)}}}const Ra=`{"title":"Use with JAX","local":"use-with-jax","sections":[{"title":"Dataset format","local":"dataset-format","sections":[{"title":"N-dimensional arrays","local":"n-dimensional-arrays","sections":[],"depth":3},{"title":"Other feature types","local":"other-feature-types","sections":[],"depth":3}],"depth":2},{"title":"Data loading","local":"data-loading","sections":[{"title":"Using with_format('jax')","local":"using-withformatjax","sections":[],"depth":3}],"depth":2}],"depth":1}`;function ka(f){return Ma(()=>{new URLSearchParams(window.location.search).get("fw")}),[]}class xa extends ya{constructor(l){super(),Ja(this,l,ka,Ia,ga,{})}}export{xa as component};
