aoc_day2_nx.exs

# part 1
File.read!('aoc/2.txt')
|> String.replace("\r", "") # i am on windows right now :(
|> String.split("\n")
|> Enum.map(fn line ->
  [direction, amt] = String.split(line)
  amt_int = String.to_integer(amt)

  case direction do
    "forward" ->
      {amt_int, 0}
    "up" ->
      {0, -amt_int}
    "down" ->
      {0, amt_int}
    _ ->
      raise "aoc is broken"
	end
end)
|> Enum.unzip()
|> then(fn {xs, ys} -> Nx.multiply(Nx.sum(Nx.tensor(xs)), Nx.sum(Nx.tensor(ys))) end)

# part 2
File.read!('aoc/2.txt')
|> String.replace("\r", "") # i am on windows right now :(
|> String.split("\n")
|> Enum.map(fn line ->
  [direction, amt] = String.split(line)
  amt_int = String.to_integer(amt)

  case direction do
    "forward" ->
      {amt_int, 0}
    "up" ->
      {0, -amt_int}
    "down" ->
      {0, amt_int}
    _ ->
      raise "aoc is broken"
	end
end)
|> Enum.unzip()
|> then(fn {forwards, aims} ->
  aims = Nx.tensor(aims)
  forwards = Nx.tensor(forwards)
  # do a cumulative sum of aims to get aims at each forward point
  n = Nx.size(aims)
  padding_config = [{n - 1, 0}]
  strides = [1]
  window_shape = {n}
  temporal_aims = Nx.window_sum(aims, window_shape, padding: padding_config)

  depths = Nx.dot(temporal_aims, forwards)
  horizontal = Nx.sum(forwards)
  Nx.multiply(depths, horizontal)
end)