Packages

exla

0.6.0
0.13.0 0.12.0 0.11.0 0.10.0 0.9.2 0.9.1 0.9.0 0.8.0 0.7.3 0.7.2 0.7.1 0.7.0 0.6.4 0.6.3 0.6.1 0.6.0 0.5.3 0.5.2 0.5.1 0.5.0 0.4.2 0.4.1 0.4.0 0.3.0 0.2.3 0.2.2 0.2.1 0.2.0

Google's XLA (Accelerated Linear Algebra) compiler/backend for Nx

HexDocs

Current section

Files

Jump to
Readme
28 Versions
4 Dependencies
Dependants
Files
Activity
exla lib exla builder.ex
Raw

lib/exla/builder.ex 

defmodule EXLA.Builder do

  @moduledoc """

  Wrapper around XLA's builder.

  """



  alias __MODULE__

  alias EXLA.{Computation, Op}



  @enforce_keys [:ref]

  defstruct [:ref, :parent, :name]



  def new(name) when is_binary(name) do

    {:ok, ref} = EXLA.NIF.new_builder(name)

    %Builder{ref: ref, parent: nil, name: name}

  end



  def new(builder = %Builder{ref: ref}, name) when is_binary(name) do

    {:ok, ref} = EXLA.NIF.create_sub_builder(ref, name)

    %Builder{ref: ref, parent: builder, name: name}

  end



  def build(root = %Op{}) do

    shape = EXLA.Op.get_shape(root)

    {:ok, ref} = EXLA.NIF.build(root.builder, root.ref)

    %Computation{ref: ref, output_shape: shape}

  end

end