zenna

module MiniOmega

using Cassette
import Base:~
export sample, unif, pointwise, <|, rt

const ID = NTuple{N, Int} where N

"Ω is a hypercube"
struct Ω

zenna

using Omega

# Beta distributed random variables in [0, 1]
p1 = beta(0.5)
pgood = beta(0.5)
pbad = beta(0.5)

# Function that simulates the car
function simulate_car(ω)
  car_is_good = Bernoulli(ω, p1(ω), Bool)

zenna

using Omega

# Add uncertainty to this if you want
region_size = 30_000 

# The number of groups
ngroups = ciid(ω -> Int(floor(uniform(ω, 1, region_size))))

# Produced random Boolean correlation matrix
function correlation_model(rng, nprobes = 10)

zenna

using Omega

# Add uncertainty to this if you want
region_size = 30_000 

# The number of groups
ngroups = uniform(1:region_size)

# Produced random Boolean correlation matrix
function correlation_model(rng, nprobes = 10)

zenna

abstract type Node end

struct Val <: Node
  x::Int
end

cons(::Val, x) = Val(x)
children(val::Val) = ()

struct Add{T1, T2} <: Node

zenna

# Simplest Version

abstract type Node end

struct Val <: Node
  x::Int
end

children(val::Val) = ()

zenna

using Omega
using Debugger
using Omega: mem

"Return sequence bounded between lb ad ub"
function f_(ω)
  lb = -1.0
  ub = 1.0
  xs = Float64[]
  x = rand(ω)

zenna

module MiniOmega
using Random

# Ω Subtypes Random.AbstractRNG to be compatible with existing samplers

mutable struct Ω <: Random.AbstractRNG
  data::Dict{Int, Any}  # Map ids to random values
  i::Int                # Current id
  logscore::Float64
end

zenna

using Arrows
using NamedTuples
import JLD: load
# import Arrows.BenchmarkArrows: STD_ROTATION_MATRIX, render
import Images: colorview, Gray

using NamedTuples
const STD_ROTATION_MATRIX = [0.94071758  -0.33430171 -0.05738258
                             -0.33835238 -0.91297877 -0.2280076
                             0.02383425  0.2339063 -0.97196698]

zenna

import ijson
import os
from enum import Enum

def clevr_iter(clevr_root,
               data_type,
               train=True):
  path = os.path.join(clevr_root, data_type)
  train_val = "train" if train else "val"
  if train: