from mxnet import gluon, nd, autograd, context
from mxnet.contrib import onnx as onnx_mxnet
import numpy as np

net = gluon.nn.HybridSequential()
net.add(gluon.nn.MaxPool1D(pool_size=2))
net.hybridize()

x = nd.random.uniform(-1,1,(1,3,4))
labels = nd.random.uniform(-1,1,(1,3,2))
L2loss = gluon.loss.L2Loss()

# showing differentiability of input
x.attach_grad()
with autograd.record():
    output = net(x)
    loss = L2loss(output, labels)

loss.backward()
print(x.grad)

# showing differentiability of output
with autograd.record():
    output = net(x)
    output.attach_grad()
    loss = L2loss(output, labels)

loss.backward()
print(output.grad)

net.export("maxpool1d")

sym = './maxpool1d-symbol.json'
params = './maxpool1d-0000.params'
onnx_file = './maxpool1d.onnx'

converted_model_path = onnx_mxnet.export_model(sym, params, [x.shape], np.float32, onnx_file)

net = gluon.nn.HybridSequential()
net.add(gluon.nn.Conv1D(1, 2, weight_initializer = 'ones'))
net.hybridize()
net.initialize(ctx=context.cpu(0))

net_params = net.collect_params('0.weight|0.bias')

x = nd.random.uniform(-1,1,(1, 1, 3))
x.attach_grad()
label = nd.random.uniform(-1,1,(1,1,2))

with autograd.record():
    output = net(x)
    loss = L2loss(output, label)

loss.backward()
print(x.grad)
print(net_params['0.weight'].grad())
print(net_params['0.bias'].grad())

with autograd.record():
    output = net(x)
    output.attach_grad()
    loss = L2loss(output, label)

loss.backward()
print(output.grad)

#export of model to onnx
net.export("conv1d")
sym = './conv1d-symbol.json'
params = './conv1d-0000.params'
onnx_file = './conv1d.onnx'

converted_model_path = onnx_mxnet.export_model(sym, params, [x.shape], np.float32, onnx_file)